#include <Wire.h>
#include <LiquidCrystal_I2C.h>
#include <DHT.h>
#include <EEPROM.h>
#include "GravityTDS.h"

#define TdsSensorPin A1
GravityTDS gravityTds;
float tdsValue = 0;


#define DHTPIN 6
#define DHTTYPE DHT22
#define RELAY_PIN_1 7
#define RELAY_PIN_2 8
#define ENCODER_PIN_A 2
#define ENCODER_PIN_B 3
#define BUTTON_PIN 4

#define FLOW_RATE 1.2
#define MAX_VOLUME 100
#define DEBOUNCE_DELAY 50
#define SENSOR_UPDATE_INTERVAL 2000
#define DISPLAY_UPDATE_INTERVAL 100

LiquidCrystal_I2C lcd(0x27, 16, 2);
DHT dht(DHTPIN, DHTTYPE);

struct SystemState {
    float humidity, temperature, lastHumidity, lastTemperature;
    int selectedPump, desiredVolume;
    bool pumpActive, displayNeedsUpdate;
    unsigned long pumpStartTime, lastSensorUpdate, lastDisplayUpdate;
    String line1, line2;
} state;

struct EncoderState {
    volatile int value, lastValue;
    unsigned long lastDebounceTime;
} encoder;

enum MenuState { IDLE, SELECT_PUMP, SET_VOLUME, CONFIRMATION, RUNNING_PUMP } currentMenu = IDLE;

void setup() {
    pinMode(RELAY_PIN_1, OUTPUT);
    pinMode(RELAY_PIN_2, OUTPUT);
    pinMode(ENCODER_PIN_A, INPUT_PULLUP);
    pinMode(ENCODER_PIN_B, INPUT_PULLUP);
    pinMode(BUTTON_PIN, INPUT_PULLUP);
    digitalWrite(RELAY_PIN_1, HIGH);
    digitalWrite(RELAY_PIN_2, HIGH);
    lcd.init();
    lcd.backlight();
    dht.begin();
    state = {0, 0, -1, -1, 1, 0, false, true, 0, 0, 0, "", ""};
    encoder = {0, 0, 0};
    attachInterrupt(digitalPinToInterrupt(ENCODER_PIN_A), encoderUpdate, CHANGE);
    
    Serial.begin(115200);
    gravityTds.setPin(TdsSensorPin);
    gravityTds.setAref(5.0);  // Reference voltage on ADC, default is 5.0V on Arduino UNO
    gravityTds.setAdcRange(1024);  // 1024 for 10bit ADC
    gravityTds.begin();  // Initialize the TDS sensor

    updateDisplay();
}

void loop() {
    unsigned long currentTime = millis();
    
    if (currentTime - state.lastSensorUpdate >= SENSOR_UPDATE_INTERVAL) {
        readSensors();
    }
    
    handleButton();
    handleMenu();
    
    if (currentMenu == RUNNING_PUMP || currentTime - state.lastDisplayUpdate >= DISPLAY_UPDATE_INTERVAL) {
        updateDisplay();
        state.lastDisplayUpdate = currentTime;
    }
}

void encoderUpdate() {
    if (millis() - encoder.lastDebounceTime > DEBOUNCE_DELAY) {
        encoder.value += (digitalRead(ENCODER_PIN_B) == LOW) ? 1 : -1;
        encoder.value = constrain(encoder.value, 0, currentMenu == SET_VOLUME ? MAX_VOLUME : (currentMenu == SELECT_PUMP || currentMenu == CONFIRMATION ? 1 : 100));
        encoder.lastDebounceTime = millis();
        state.displayNeedsUpdate = true;
    }
}

void readSensors() {
    float newHumidity = dht.readHumidity(), newTemperature = dht.readTemperature();

    gravityTds.setTemperature(newTemperature);  // Use temperature from DHT for compensation
    gravityTds.update();  // Sample and calculate the TDS value
    tdsValue = gravityTds.getTdsValue();  // Get TDS value in ppm


    if (!isnan(newHumidity) && !isnan(newTemperature) && (newHumidity != state.lastHumidity || newTemperature != state.lastTemperature)) {
        state.humidity = newHumidity;
        state.temperature = newTemperature;
        state.lastHumidity = newHumidity;
        state.lastTemperature = newTemperature;
        if (currentMenu == IDLE) state.displayNeedsUpdate = true;
    }
    state.lastSensorUpdate = millis();
}

void updateDisplay() {
    String newLine1, newLine2;
    unsigned long currentTime = millis();
    
    switch (currentMenu) {
        case IDLE:
            newLine1 = "H:" + String(state.humidity, 1) + "% P:" + String(tdsValue, 0);
            newLine2 = "T:" + String(state.temperature, 1) + "C";
            break;
        case SELECT_PUMP:
            newLine1 = "Select Pump:";
            newLine2 = "Pump " + String((encoder.value % 2) + 1);
            break;
        case SET_VOLUME:
            newLine1 = "Set Vol. P" + String(state.selectedPump);
            newLine2 = String(state.desiredVolume) + " mL";
            break;
        case CONFIRMATION:
            newLine1 = "P" + String(state.selectedPump) + " " + String(state.desiredVolume) + "mL OK?";
            newLine2 = (encoder.value % 2 == 0) ? "Yes" : "No";
            break;
        case RUNNING_PUMP:
            newLine1 = "Pump " + String(state.selectedPump) + " Run";
            unsigned long elapsedTime = (currentTime - state.pumpStartTime) / 1000;
            unsigned long totalTime = (unsigned long)(state.desiredVolume / FLOW_RATE);
            newLine2 = String(elapsedTime) + "s/" + String(totalTime) + "s";
            break;
    }
    
    if (newLine1 != state.line1 || newLine2 != state.line2) {
        lcd.clear();
        lcd.setCursor(0, 0);
        lcd.print(newLine1);
        lcd.setCursor(0, 1);
        lcd.print(newLine2);
        state.line1 = newLine1;
        state.line2 = newLine2;
    }
}

void handleButton() {
    static unsigned long lastButtonPress = 0;
    static bool buttonWasPressed = false;
    
    if (digitalRead(BUTTON_PIN) == LOW && !buttonWasPressed && (millis() - lastButtonPress > 250)) {
        buttonWasPressed = true;
        lastButtonPress = millis();
        
        switch (currentMenu) {
            case IDLE:
                currentMenu = SELECT_PUMP;
                encoder.value = state.selectedPump - 1;
                break;
            case SELECT_PUMP:
                state.selectedPump = (encoder.value % 2) + 1;
                currentMenu = SET_VOLUME;
                encoder.value = state.desiredVolume;
                break;
            case SET_VOLUME:
                if (state.desiredVolume > 0) {
                    currentMenu = CONFIRMATION;
                    encoder.value = 0;
                }
                break;
            case CONFIRMATION:
                if (encoder.value % 2 == 0) {
                    startPump();
                    currentMenu = RUNNING_PUMP;
                } else {
                    currentMenu = SET_VOLUME;
                    encoder.value = state.desiredVolume;
                }
                break;
        }
        state.displayNeedsUpdate = true;
    } else if (digitalRead(BUTTON_PIN) == HIGH) {
        buttonWasPressed = false;
    }
}

void handleMenu() {
    if (encoder.value != encoder.lastValue) {
        if (currentMenu == SET_VOLUME) state.desiredVolume = encoder.value;
        encoder.lastValue = encoder.value;
        state.displayNeedsUpdate = true;
    }
    
    if (currentMenu == RUNNING_PUMP && state.pumpActive) {
        unsigned long pumpRunTime = (unsigned long)(state.desiredVolume / FLOW_RATE) * 1000;
        if (millis() - state.pumpStartTime >= pumpRunTime) {
            stopPump();
            currentMenu = IDLE;
            state.displayNeedsUpdate = true;
        }
    }
}

void startPump() {
    digitalWrite(state.selectedPump == 1 ? RELAY_PIN_1 : RELAY_PIN_2, LOW);
    state.pumpActive = true;
    state.pumpStartTime = millis();
    state.displayNeedsUpdate = true;
}

void stopPump() {
    digitalWrite(RELAY_PIN_1, HIGH);
    digitalWrite(RELAY_PIN_2, HIGH);
    state.pumpActive = false;
    state.displayNeedsUpdate = true;
}