#include <Arduino.h>
#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <RTClib.h>

static const uint8_t CHANNELS = 4;
static const uint8_t PIN_PUMP[CHANNELS] = {16, 17, 18, 19};
static const uint8_t PIN_BUTTON_UP = 32;
static const uint8_t PIN_BUTTON_DOWN = 33;
static const uint8_t PIN_BUTTON_OK = 25;
static const uint8_t PIN_BUZZER = 26;
static const uint8_t PIN_LED_POWER = 27;
static const uint8_t PIN_LED_FAULT = 14;
static const uint16_t DEMO_RUN_SEC = 5;
static const uint16_t SAFETY_LIMIT_SEC = 6;

LiquidCrystal_I2C lcd(0x27, 16, 2);
RTC_DS3231 rtc;

uint8_t selectedChannel = 0;
bool pumpRunning[CHANNELS] = {false, false, false, false};
uint32_t pumpStartedMs[CHANNELS] = {0, 0, 0, 0};
bool fault = false;
uint32_t lastButtonMs = 0;
uint32_t lastLcdMs = 0;

void stopAllPumps() {
  for (uint8_t ch = 0; ch < CHANNELS; ch++) {
    digitalWrite(PIN_PUMP[ch], LOW);
    pumpRunning[ch] = false;
  }
}

void setFault(const char *message) {
  stopAllPumps();
  fault = true;
  digitalWrite(PIN_LED_FAULT, HIGH);
  digitalWrite(PIN_BUZZER, HIGH);
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("SAFETY FAULT");
  lcd.setCursor(0, 1);
  lcd.print(message);
}

void clearFault() {
  fault = false;
  digitalWrite(PIN_LED_FAULT, LOW);
  digitalWrite(PIN_BUZZER, LOW);
}

void startDose(uint8_t ch) {
  clearFault();
  digitalWrite(PIN_PUMP[ch], HIGH);
  pumpStartedMs[ch] = millis();
  pumpRunning[ch] = true;
}

void updatePumps() {
  uint32_t now = millis();
  for (uint8_t ch = 0; ch < CHANNELS; ch++) {
    if (!pumpRunning[ch]) continue;
    uint32_t elapsedSec = (now - pumpStartedMs[ch]) / 1000UL;
    if (elapsedSec > SAFETY_LIMIT_SEC) {
      setFault("TIMEOUT");
      return;
    }
    if (elapsedSec >= DEMO_RUN_SEC) {
      digitalWrite(PIN_PUMP[ch], LOW);
      pumpRunning[ch] = false;
    }
  }
}

void updateButtons() {
  if (millis() - lastButtonMs < 250) return;

  bool up = digitalRead(PIN_BUTTON_UP) == LOW;
  bool down = digitalRead(PIN_BUTTON_DOWN) == LOW;
  bool ok = digitalRead(PIN_BUTTON_OK) == LOW;

  if (up && down) {
    lastButtonMs = millis();
    setFault("BTN TEST");
    return;
  }

  if (up) {
    selectedChannel = (selectedChannel + 1) % CHANNELS;
    lastButtonMs = millis();
  } else if (down) {
    selectedChannel = selectedChannel == 0 ? CHANNELS - 1 : selectedChannel - 1;
    lastButtonMs = millis();
  } else if (ok) {
    startDose(selectedChannel);
    lastButtonMs = millis();
  }
}

void updateLcd() {
  if (fault || millis() - lastLcdMs < 500) return;
  lastLcdMs = millis();

  DateTime now = rtc.now();
  lcd.setCursor(0, 0);
  char top[17];
  snprintf(top, sizeof(top), "%02d:%02d CH%d SEL", now.hour(), now.minute(), selectedChannel + 1);
  lcd.print(top);
  for (uint8_t i = strlen(top); i < 16; i++) lcd.print(' ');

  lcd.setCursor(0, 1);
  int active = -1;
  for (uint8_t ch = 0; ch < CHANNELS; ch++) {
    if (pumpRunning[ch]) active = ch;
  }
  if (active >= 0) {
    uint32_t elapsedSec = (millis() - pumpStartedMs[active]) / 1000UL;
    char line[17];
    snprintf(line, sizeof(line), "CH%d DOSING %lus", active + 1, (unsigned long)elapsedSec);
    lcd.print(line);
    for (uint8_t i = strlen(line); i < 16; i++) lcd.print(' ');
  } else {
    lcd.print("READY           ");
  }
}

void setup() {
  for (uint8_t ch = 0; ch < CHANNELS; ch++) {
    pinMode(PIN_PUMP[ch], OUTPUT);
    digitalWrite(PIN_PUMP[ch], LOW);
  }
  pinMode(PIN_BUTTON_UP, INPUT_PULLUP);
  pinMode(PIN_BUTTON_DOWN, INPUT_PULLUP);
  pinMode(PIN_BUTTON_OK, INPUT_PULLUP);
  pinMode(PIN_BUZZER, OUTPUT);
  pinMode(PIN_LED_POWER, OUTPUT);
  pinMode(PIN_LED_FAULT, OUTPUT);
  digitalWrite(PIN_BUZZER, LOW);
  digitalWrite(PIN_LED_FAULT, LOW);
  digitalWrite(PIN_LED_POWER, HIGH);

  Wire.begin(21, 22);
  lcd.init();
  lcd.backlight();
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("AquaDose Demo");
  lcd.setCursor(0, 1);
  lcd.print("All pumps OFF");

  rtc.begin();
  if (rtc.lostPower()) {
    rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));
  }
  delay(1200);
}

void loop() {
  updateButtons();
  updatePumps();
  updateLcd();
}