#include <U8g2lib.h>
#include "HX711.h"

/* ================= ПИНЫ ================= */
#define PIN_PUMP_RELAY       26
#define PIN_POWER_RELAY      25
#define PIN_BUTTON           27
#define BUZZER_PIN           33

#define PIN_HX711_DT         16
#define PIN_HX711_SCK        4

/* ================= ВЕСА ================= */
#define EMPTY_WEIGHT         1060
#define MIN_WEIGHT           1560    // EMPTY_WEIGHT + 500 мл
#define FULL_WEIGHT          2760    // EMPTY_WEIGHT + 1700 мл
#define CUP_VOLUME           250
#define WEIGHT_HYST          20      // вес воды для гистерезиса

#define NO_FLOW_THRESHOLD_MS 5000    // 5 секунд для детекции "нет потока"
#define FILL_TIMEOUT_MS      120000  // 2 минуты на максимальный налив

/* ================= КНОПКА ================= */
#define DEBOUNCE_DELAY    50
#define LONG_PRESS_MS     3000
#define DOUBLE_CLICK_MS   400

/* ================= ДИСПЛЕЙ ================= */
U8G2_SSD1306_128X64_NONAME_F_HW_I2C display(U8G2_R0, U8X8_PIN_NONE);

/* ================= ИКОНКА ПИТАНИЯ 16х16 ================= */
const unsigned char icon_power_16x16 [] PROGMEM = {
  0x80,0x01,0x80,0x01,0x80,0x01,0x98,0x19,
  0x8c,0x31,0x86,0x61,0x82,0x41,0x83,0xc1,
  0x03,0xc0,0x03,0xc0,0x03,0xc0,0x02,0x40,
  0x06,0x60,0x0c,0x30,0x78,0x1e,0xe0,0x07
};

/* ================= ИКОНКА Wi-Fi 16х16 ================= */
const unsigned char icon_wifi_16x16 [] PROGMEM = {
  0x00, 0x00, 0x00, 0x00,
  0xF0, 0x0F, 0x3C, 0x3C,
  0x06, 0x60, 0xE3, 0xC7,
  0xF8, 0x1F, 0x18, 0x18,
  0xC0, 0x03, 0xE0, 0x07,
  0x00, 0x00, 0x80, 0x01,
  0x80, 0x01, 0x80, 0x01,
  0x00, 0x00, 0x00, 0x00
};

/* ================= ИКОНКА ЧАШКИ 20х20 ================= */
const unsigned char cup_20x20 [] PROGMEM = {
  0xFE, 0x3F, 0xF0, 0xFF, 0x7F, 0xF0,
  0x03, 0x60, 0xF0, 0x03, 0xE0, 0xF3,
  0x03, 0x60, 0xF7, 0x03, 0x60, 0xF4,
  0x03, 0x60, 0xFC, 0x03, 0x60, 0xFC,
  0x03, 0x60, 0xFC, 0x03, 0x60, 0xFC,
  0x03, 0x60, 0xFC, 0x03, 0x60, 0xFC,
  0x03, 0x60, 0xFC, 0x03, 0x60, 0xF6,
  0x03, 0xE0, 0xF7, 0x03, 0xE0, 0xF3,
  0x07, 0x70, 0xF0, 0xFE, 0x3F, 0xF0,
  0xFC, 0x1F, 0xF0, 0x00, 0x00, 0xF0
};

/* ================= СОСТОЯНИЯ ================= */
enum SystemState { 
  ST_INIT,      // Инициализация
  ST_IDLE,      // Ожидание
  ST_FILLING,   // Идет налив
  ST_ERROR      // Аварийная остановка
};

enum ErrorCode { 
  ERR_NONE, 
  ERR_HX711_TIMEOUT,    // Ошибка весов
  ERR_NO_FLOW,          // Нет потока воды
  ERR_FILL_TIMEOUT      // Превышено время налива
};

SystemState currentState = ST_INIT;
ErrorCode lastError = ERR_NONE;

/* ================= ГЛОБАЛЬНЫЕ ================= */
HX711 scale;

float currentWeight = 0;
float targetWeight  = 0;
float fillStartVolume = 0; // Объем воды в начале налива

bool kettlePresent = false;
bool powerRelayState = false;
bool powerRelayWasSet = false; // Для предотвращения дребезга

unsigned long fillStartTime = 0;
unsigned long lastWeightChangeTime = 0;
unsigned long lastPowerRelayUpdate = 0;

/* ================= КНОПКА ================= */
bool buttonPressed = false;
unsigned long buttonPressTime = 0;
unsigned long lastClickTime = 0;
uint8_t buttonClickCount = 0;

/* ================= УТИЛИТЫ ================= */
// Функция для преобразования мл в кружки (1 кружка = 250 мл)
// Остаток отбрасывается: 550 мл = 2 кружки, 950 мл = 3 кружки
int mlToCups(float ml) {
  if (ml <= 0) return 0;
  int cups = ml / CUP_VOLUME; // Целочисленное деление - остаток отбрасывается
  return cups;
}

// Функция для форматирования строки с количеством кружек (только число)
String formatCupsNumber(int cups) {
  return String(cups);
}

/* ================= БУЗЗЕР ================= */
void buzzerShort(uint8_t n) {
  for (uint8_t i = 0; i < n; i++) {
    digitalWrite(BUZZER_PIN, HIGH); 
    delay(80);
    digitalWrite(BUZZER_PIN, LOW);  
    delay(120);
  }
}

void buzzerLong(uint8_t n) {
  for (uint8_t i = 0; i < n; i++) {
    digitalWrite(BUZZER_PIN, HIGH); 
    delay(1000);
    digitalWrite(BUZZER_PIN, LOW);  
    delay(200);
  }
}

/* ================= РЕЛЕ ================= */
void setPumpRelay(bool on)  { 
  digitalWrite(PIN_PUMP_RELAY, on ? HIGH : LOW); 
}

void setPowerRelay(bool on) { 
  // Защита от дребезга - меняем состояние только если оно изменилось
  if (on != powerRelayWasSet) {
    digitalWrite(PIN_POWER_RELAY, on ? HIGH : LOW);
    powerRelayWasSet = on;
    powerRelayState = on; // Обновляем глобальную переменную
    Serial.print("Power relay: ");
    Serial.println(on ? "ON" : "OFF");
  }
}

/* ================= УТИЛИТЫ ================= */
bool isKettlePresent() {
  return currentWeight >= (EMPTY_WEIGHT - WEIGHT_HYST);
}

bool isMinWaterReached() {
  return currentWeight >= (MIN_WEIGHT - WEIGHT_HYST);
}

// Функция расчета объема воды в мл
float getWaterVolume() {
  float vol = currentWeight - EMPTY_WEIGHT;
  return (vol < 0) ? 0 : vol;
}

// Функция расчета целевого объема воды в мл
float getTargetWaterVolume() {
  float vol = targetWeight - EMPTY_WEIGHT;
  return (vol < 0) ? 0 : vol;
}

/* ================= ОШИБКА ================= */
void setError(ErrorCode err) {
  currentState = ST_ERROR;
  lastError = err;
  setPumpRelay(false);
  setPowerRelay(false);
  buzzerLong(3);
  
  Serial.print("ERROR: ");
  switch (err) {
    case ERR_HX711_TIMEOUT: Serial.println("Weight sensor timeout"); break;
    case ERR_NO_FLOW: Serial.println("No water flow"); break;
    case ERR_FILL_TIMEOUT: Serial.println("Filling timeout"); break;
    default: break;
  }
}

/* ================= ВЕС ================= */
void updateWeight() {
  static unsigned long lastRead = 0;
  static unsigned long hxTimeout = 0;

  if (millis() - lastRead < 200) return;
  lastRead = millis();

  if (scale.is_ready()) {
    hxTimeout = 0;
    float w = scale.get_units(1);
    if (abs(w - currentWeight) > 0.1) {
      currentWeight = w;
      lastWeightChangeTime = millis();
    }
  } else {
    if (hxTimeout == 0) hxTimeout = millis();
    if (millis() - hxTimeout > 2000) {
      setError(ERR_HX711_TIMEOUT);
    }
  }
}

/* ================= УПРАВЛЕНИЕ ПИТАНИЕМ ЧАЙНИКА ================= */
void updatePowerRelayControl() {
  // Обновляем не чаще чем раз в 2 секунды для предотвращения дребезга
  if (millis() - lastPowerRelayUpdate < 2000) return;
  lastPowerRelayUpdate = millis();
  
  bool shouldBeOn = isMinWaterReached();
  setPowerRelay(shouldBeOn);
}

/* ================= КНОПКА ================= */
void handleButton() {
  // В режиме ошибки кнопка полностью отключена
  if (currentState == ST_ERROR) {
    buttonPressed = false;
    buttonClickCount = 0;
    return;
  }
  
  // Кнопка активна только при наличии чайника
  if (!kettlePresent) {
    buttonPressed = false;
    buttonClickCount = 0;
    return;
  }

  unsigned long now = millis();
  bool pressed = digitalRead(PIN_BUTTON) == LOW;

  /* ===== РЕЖИМ НАЛИВА: ТОЛЬКО LONG PRESS ===== */
  if (currentState == ST_FILLING) {
    if (pressed && !buttonPressed) {
      buttonPressed = true;
      buttonPressTime = now;
    }
    if (!pressed && buttonPressed) {
      buttonPressed = false;
      if (now - buttonPressTime >= LONG_PRESS_MS) {
        setPumpRelay(false);
        currentState = ST_IDLE;
        buzzerShort(3);
        Serial.println("Manual stop (long press)");
      }
    }
    return;
  }

  /* ===== IDLE: ОБРАБОТКА КЛИКОВ ===== */
  if (pressed && !buttonPressed) {
    buttonPressed = true;
    buttonPressTime = now;
  }

  if (!pressed && buttonPressed) {
    buttonPressed = false;
    unsigned long dt = now - buttonPressTime;
    
    // Долгое нажатие - уже обработано в режиме FILLING
    if (dt >= LONG_PRESS_MS) return;

    // Короткое нажатие - обработка кликов
    if (dt > DEBOUNCE_DELAY) {
      if (now - lastClickTime < DOUBLE_CLICK_MS) {
        buttonClickCount++;
      } else {
        buttonClickCount = 1;
      }
      lastClickTime = now;
    }
  }

  // Выполнение действия после таймаута двойного клика
  if (buttonClickCount > 0 && millis() - lastClickTime > DOUBLE_CLICK_MS) {
    float add = 0;

    if (buttonClickCount == 1) {
      // Одинарный клик
      if (currentWeight < MIN_WEIGHT - WEIGHT_HYST) {
        add = MIN_WEIGHT - currentWeight; // Налить до минимального (500 мл = 2 кружки)
      } else {
        add = CUP_VOLUME; // Налить одну чашку (250 мл = 1 кружка)
      }
    } else if (buttonClickCount == 2) {
      // Двойной клик - налить до полного
      add = FULL_WEIGHT - currentWeight;
    }

    // Проверка на корректность объема
    if (add > WEIGHT_HYST && currentWeight + add <= FULL_WEIGHT + WEIGHT_HYST) {
      targetWeight = currentWeight + add;
      fillStartVolume = getWaterVolume(); // Запоминаем объем в начале налива
      currentState = ST_FILLING;
      fillStartTime = millis();
      lastWeightChangeTime = millis();
      setPumpRelay(true);
      buzzerShort(1);
      
      Serial.print("Start filling: ");
      Serial.print(add, 0);
      Serial.print("g to ");
      Serial.print(targetWeight, 0);
      Serial.println("g");
      Serial.print("Start volume: ");
      Serial.print(fillStartVolume);
      Serial.print(" ml (");
      Serial.print(mlToCups(fillStartVolume));
      Serial.print(" cups), Target volume: ");
      Serial.print(getTargetWaterVolume());
      Serial.print(" ml (");
      Serial.print(mlToCups(getTargetWaterVolume()));
      Serial.println(" cups)");
    }

    buttonClickCount = 0;
  }
}

/* ================= НАЛИВ ================= */
void handleFilling() {
  // Проверка таймаута налива (2 минуты)
  if (millis() - fillStartTime > FILL_TIMEOUT_MS) {
    setError(ERR_FILL_TIMEOUT);
    return;
  }

  // Обновляем время последнего изменения веса если вес изменился
  float currentWater = getWaterVolume();
  static float lastWater = 0;
  
  if (abs(currentWater - lastWater) > 5) { // Порог 5г для NO_FLOW
    lastWater = currentWater;
    lastWeightChangeTime = millis();
  }

  // Проверка отсутствия потока воды
  if (millis() - lastWeightChangeTime > NO_FLOW_THRESHOLD_MS &&
      currentWeight < targetWeight - 50) { // 50г порог для срабатывания
    setError(ERR_NO_FLOW);
    return;
  }

  // Проверка достижения целевого веса
  if (currentWeight >= targetWeight - WEIGHT_HYST) {
    setPumpRelay(false);
    currentState = ST_IDLE;
    buzzerShort(2);
    Serial.println("Filling completed");
  }
}

/* ================= ДИСПЛЕЙ ================= */
void drawProgressBar(int x, int y, int w, int h, int p){
  display.drawFrame(x, y, w, h);
  display.drawBox(x + 1, y + 1, (w - 2) * p / 100, h - 2);
}

void updateDisplay() {
  display.clearBuffer();
  
  if (currentState == ST_INIT) {
    // Экран загрузки
    display.setFont(u8g2_font_10x20_tf);
    display.drawStr(15, 20, "SMART PUMP");
    display.drawStr(25, 45, "LOADING...");
  }
  else if (currentState == ST_ERROR) {
    // Режим ошибки
    display.setFont(u8g2_font_6x10_tf);
    display.setCursor(0, 0);
    display.print("ERROR:");
    
    display.setCursor(0, 12);
    switch (lastError) {
      case ERR_HX711_TIMEOUT: 
        display.print("Weight sensor fail"); 
        break;
      case ERR_NO_FLOW: 
        display.print("No water flow!"); 
        break;
      case ERR_FILL_TIMEOUT: 
        display.print("Filling timeout"); 
        break;
      default: 
        display.print("Unknown error"); 
        break;
    }
    
    display.setCursor(0, 40);
    display.print("Reboot required");
  }
  else if (currentState == ST_IDLE) {
    // Режим ожидания
    
    // Иконка питания в левом верхнем углу (ВСЕГДА когда powerRelayState = true)
    if (powerRelayState) {
      display.drawXBMP(0, 0, 16, 16, icon_power_16x16);
    }
    
    // Иконка Wi-Fi в правом верхнем углу (всегда показываем)
    display.drawXBMP(112, 0, 16, 16, icon_wifi_16x16); // 128 - 16 = 112
    
    // Статус работы помпы в верхней строке по центру
    display.setFont(u8g2_font_6x10_tf);
    String statusStr;
    if (!kettlePresent) {
      statusStr = "NO KETTLE";
    } else {
      statusStr = "Ready";
    }
    int statusWidth = display.getStrWidth(statusStr.c_str());
    display.setCursor((128 - statusWidth) / 2, 0);
    display.print(statusStr);
    
    // Объем воды в кружках по центру экрана
    display.setFont(u8g2_font_10x20_tf);
    float waterVol = getWaterVolume();
    int cups = mlToCups(waterVol);
    String cupsStr = formatCupsNumber(cups);
    
    // Рассчитываем ширину цифры
    int cupsWidth = display.getStrWidth(cupsStr.c_str());
    
    // Центрируем композицию: цифра + иконка 20px + отступ
    int totalWidth = cupsWidth + 20 + 4; // цифра + иконка + отступ
    int startX = (128 - totalWidth) / 2;
    
    // Выравниваем вертикально: текст высотой 20px, иконка 20px
    int textY = 28; // Позиция текста
    int iconY = 24; // Позиция иконки
    
    // Рисуем число
    display.setCursor(startX, textY);
    display.print(cupsStr);
    
    // Рисуем иконку чашки 20x20 справа от числа
    display.drawXBMP(startX + cupsWidth + 4, iconY, 20, 20, cup_20x20);
  }
  else if (currentState == ST_FILLING) {
    // Режим налива
    // Иконка питания в левом верхнем углу (ВСЕГДА когда powerRelayState = true)
    if (powerRelayState) {
      display.drawXBMP(0, 0, 16, 16, icon_power_16x16);
    }
    
    // Иконка Wi-Fi в правом верхнем углу (всегда показываем)
    display.drawXBMP(112, 0, 16, 16, icon_wifi_16x16); // 128 - 16 = 112
    
    // Статус работы помпы в верхней строке по центру
    display.setFont(u8g2_font_6x10_tf);
    String statusStr = "Filling...";
    int statusWidth = display.getStrWidth(statusStr.c_str());
    display.setCursor((128 - statusWidth) / 2, 0);
    display.print(statusStr);
    
    // Текущие кружки -> Целевые кружки
    display.setFont(u8g2_font_10x20_tf);
    
    // Рассчитываем текущее и целевое количество кружек
    float currentWater = getWaterVolume();
    float targetWaterVolume = getTargetWaterVolume();
    
    int currentCups = mlToCups(currentWater);
    int targetCups = mlToCups(targetWaterVolume);
    
    // Форматируем строку: "текущее -> целевое"
    String cupsStr = formatCupsNumber(currentCups) + " -> " + formatCupsNumber(targetCups);
    int cupsWidth = display.getStrWidth(cupsStr.c_str());
    
    // Центрируем: текст + иконка 20px + отступ
    int totalWidth = cupsWidth + 20 + 4;
    int startX = (128 - totalWidth) / 2;
    
    // Вертикальное выравнивание
    int textY = 13; // Выше для режима FILLING
    int iconY = 8;  // Соответственно выше для иконки
    
    // Рисуем текст
    display.setCursor(startX, textY);
    display.print(cupsStr);
    
    // Рисуем иконку чашки 20x20 справа от текста
    display.drawXBMP(startX + cupsWidth + 4, iconY, 20, 20, cup_20x20);
    
    // Прогресс-бар
    int progress = 0;
    if (targetWaterVolume > fillStartVolume) {
      // Расчет прогресса: от начального объема до целевого (в мл для точности)
      progress = map(currentWater, fillStartVolume, targetWaterVolume, 0, 100);
      progress = constrain(progress, 0, 100);
    }
    
    // Проценты по центру экрана
    display.setFont(u8g2_font_10x20_tf);
    String percentStr = String(progress) + "%";
    int percentWidth = display.getStrWidth(percentStr.c_str());
    display.setCursor((128 - percentWidth) / 2, 35);
    display.print(percentStr);
    
    // Прогресс-бар под процентами
    drawProgressBar(14, 50, 100, 8, progress);
    
    // Отладочная информация в Serial
    static unsigned long lastDebug = 0;
    if (millis() - lastDebug > 1000) {
      Serial.print("Filling: ");
      Serial.print(currentWater);
      Serial.print(" ml (");
      Serial.print(currentCups);
      Serial.print(" cups) -> ");
      Serial.print(targetWaterVolume);
      Serial.print(" ml (");
      Serial.print(targetCups);
      Serial.print(" cups) = ");
      Serial.print(progress);
      Serial.println("%");
      lastDebug = millis();
    }
  }
  
  display.sendBuffer();
}

/* ================= SETUP ================= */
void setup() {
  Serial.begin(115200);
  Serial.println("Smart Kettle Pump - Starting...");
  
  // Инициализация пинов
  pinMode(PIN_PUMP_RELAY, OUTPUT);
  pinMode(PIN_POWER_RELAY, OUTPUT);
  pinMode(PIN_BUTTON, INPUT_PULLUP);
  pinMode(BUZZER_PIN, OUTPUT);
  
  // Выключить все реле
  setPumpRelay(false);
  setPowerRelay(false);
  
  // Инициализация дисплея
  display.begin();
  display.setFont(u8g2_font_6x10_tf);
  display.setFontRefHeightExtendedText();
  display.setDrawColor(1);
  display.setFontPosTop();
  display.setFontDirection(0);
  
  // Показываем экран загрузки
  display.clearBuffer();
  display.setFont(u8g2_font_10x20_tf);
  display.drawStr(15, 20, "SMART PUMP");
  display.drawStr(25, 45, "LOADING...");
  display.sendBuffer();
  
  // Инициализация весов
  scale.begin(PIN_HX711_DT, PIN_HX711_SCK);
  scale.set_scale(0.42); // Калибровочный коэффициент
  
  // Ждем инициализацию HX711
  unsigned long startTime = millis();
  while (!scale.is_ready() && (millis() - startTime < 3000)) {
    delay(10);
  }
  
  if (scale.is_ready()) {
    // Читаем начальный вес
    currentWeight = scale.get_units(5);
    currentState = ST_IDLE;
    buzzerShort(1);
    
    // Инициализируем статусы
    kettlePresent = isKettlePresent();
    powerRelayState = isMinWaterReached();
    setPowerRelay(powerRelayState); // Устанавливаем начальное состояние
    
    Serial.println("System initialized successfully");
    Serial.print("Initial weight: ");
    Serial.print(currentWeight, 1);
    Serial.print(" g, Water: ");
    Serial.print(getWaterVolume());
    Serial.print(" ml (");
    Serial.print(mlToCups(getWaterVolume()));
    Serial.println(" cups)");
    Serial.print("Kettle present: ");
    Serial.println(kettlePresent ? "YES" : "NO");
    Serial.print("Power relay: ");
    Serial.println(powerRelayState ? "ON" : "OFF");
  } else {
    setError(ERR_HX711_TIMEOUT);
  }
}

/* ================= LOOP ================= */
void loop() {
  // Обновляем вес
  updateWeight();
  
  // Определяем наличие чайника
  kettlePresent = isKettlePresent();
  
  // Управление реле питания чайника (с защитой от дребезга)
  updatePowerRelayControl();
  
  // Обработка ошибки - блокирующее состояние
  if (currentState == ST_ERROR) {
    // Мигаем дисплеем и пищим периодически (1 короткий каждые 5 сек)
    static unsigned long lastErrorBeep = 0;
    if (millis() - lastErrorBeep > 5000) {
      buzzerShort(1);
      lastErrorBeep = millis();
    }
    updateDisplay();
    delay(100);
    return;
  }
  
  // Показываем экран загрузки если еще в инициализации
  if (currentState == ST_INIT) {
    updateDisplay();
    delay(100);
    return;
  }
  
  // Обработка кнопки
  handleButton();
  
  // Обработка налива если он активен
  if (currentState == ST_FILLING) {
    handleFilling();
  }
  
  // Обновление дисплея
  updateDisplay();
  
  // Небольшая задержка для стабильности
  delay(50);
}