#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <RTClib.h>
#include <TimeLib.h>

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define SCREEN_ADDRESS 0x3C

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
RTC_DS3231 rtc;

const unsigned long DEBOUNCE_DELAY = 1000;
volatile unsigned long lastDebounceTime1 = 0;
volatile unsigned long lastDebounceTime2 = 0;
volatile bool button1Pressed = false;
volatile bool button2Pressed = false;

class GrowSystem {
public:
    DateTime endTime;
    TimeSpan totalDuration;
    TimeSpan stage1Duration, stage2Duration, stage3Duration;
    TimeSpan ledOnDuration, ledOffDuration;
    bool countdownStarted, countdownOver, ledOn;
    int currentStage;
    DateTime ledStartTime;
    int startButtonPin, ledPin, fanPin;
    String itemName, floorName;

    GrowSystem(TimeSpan total, TimeSpan stage1, TimeSpan stage2, TimeSpan stage3, 
               TimeSpan ledOn, TimeSpan ledOff, int startPin, int ledPinNum, int fanPinNum, 
               String item, String floor)
    : totalDuration(total), stage1Duration(stage1), stage2Duration(stage2), stage3Duration(stage3),
      ledOnDuration(ledOn), ledOffDuration(ledOff), startButtonPin(startPin), ledPin(ledPinNum), 
      fanPin(fanPinNum), itemName(item), floorName(floor), lastPrintedTime(0) {
        if ((stage1Duration + stage2Duration + stage3Duration).totalseconds() != totalDuration.totalseconds()) {
            Serial.println("Warning: Sum of stage durations does not match total duration!");
        }
        reset();
    }

    void start() {
        countdownStarted = true;
        countdownOver = false;
        currentStage = 1;
        endTime = rtc.now() + totalDuration;
        ledStartTime = rtc.now();
        ledOn = true;
        digitalWrite(ledPin, HIGH);
        digitalWrite(fanPin, HIGH);
        Serial.println(floorName + ": Countdown started! LED and fan turned on");
    }

    void update() {
        if (countdownStarted && !countdownOver) {
            DateTime now = rtc.now();
            TimeSpan elapsed = now - (endTime - totalDuration);
            TimeSpan remaining = endTime - now;

            if (remaining.totalseconds() <= 0) {
                finish();
            } else {
                updateStage(elapsed);
                updateLED();
            }
        }
    }

    void updateDisplay(int yOffset) {
        display.setCursor(0, yOffset);
        display.print(floorName + " : ");
        
        if (countdownStarted && !countdownOver) {
            display.println(itemName);
            TimeSpan timeLeft = TimeSpan(endTime.unixtime() - rtc.now().unixtime());
            display.setCursor(0, yOffset + 10);
            display.print("Stage ");
            display.print(currentStage);
            display.print(": ");
            
            if (timeLeft.days() > 0) {
                display.print(timeLeft.days());
                display.print("D ");
            }
            if (timeLeft.hours() > 0 || timeLeft.days() > 0) {
                display.print(timeLeft.hours());
                display.print("H ");
            }
            display.print(timeLeft.minutes());
            display.print("M ");
            display.print(timeLeft.seconds());
            display.print("S");
            
            if (timeLeft.totalseconds() != lastPrintedTime.totalseconds()) {
                printFullTimeToSerial(timeLeft);
                lastPrintedTime = timeLeft;
            }
        } else if (countdownOver) {
            display.println(itemName);
            display.setCursor(0, yOffset + 10);
            display.println("HARVEST!");
        } else {
            display.println("Empty");
            display.setCursor(0, yOffset + 10);
            display.println("Press start");
        }
    }

    void reset() {
        countdownStarted = false;
        countdownOver = false;
        currentStage = 0;
        ledOn = false;
        digitalWrite(ledPin, LOW);
        digitalWrite(fanPin, LOW);
        lastPrintedTime = TimeSpan(0);
    }

private:
    TimeSpan lastPrintedTime;
    
    void finish() {
        countdownStarted = false;
        countdownOver = true;
        digitalWrite(ledPin, LOW);
        digitalWrite(fanPin, LOW);
        Serial.println(floorName + ": Countdown finished. Time to harvest! LED and fan turned off");
        lastPrintedTime = TimeSpan(0);
    }

    void updateStage(TimeSpan elapsed) {
        if (elapsed.totalseconds() >= (stage1Duration + stage2Duration).totalseconds()) {
            if (currentStage != 3) {
                currentStage = 3;
                Serial.println(floorName + ": Entered Stage 3");
            }
        } else if (elapsed.totalseconds() >= stage1Duration.totalseconds()) {
            if (currentStage != 2) {
                currentStage = 2;
                Serial.println(floorName + ": Entered Stage 2");
            }
        }
    }

    void updateLED() {
        TimeSpan elapsedTime = rtc.now() - ledStartTime;
        if ((ledOn && elapsedTime.totalseconds() >= ledOnDuration.totalseconds()) || 
            (!ledOn && elapsedTime.totalseconds() >= ledOffDuration.totalseconds())) {
            ledOn = !ledOn;
            digitalWrite(ledPin, ledOn ? HIGH : LOW);
            ledStartTime = rtc.now();
            Serial.println(floorName + ": LED turned " + String(ledOn ? "on" : "off"));
        }
    }

    void printFullTimeToSerial(TimeSpan timeLeft) {
        Serial.printf("%s: Stage %d - %d DAY %d HR %d MIN %d SEC\n", 
                      floorName.c_str(), currentStage, timeLeft.days(), timeLeft.hours(), 
                      timeLeft.minutes(), timeLeft.seconds());
    }
};

GrowSystem system1(TimeSpan(0, 0, 3, 0),  // total: 45 วัน
                   TimeSpan(0, 0, 1, 0),   // stage 1: 7 วัน
                   TimeSpan(0, 0, 1, 0),  // stage 2: 30 วัน
                   TimeSpan(0, 0, 1, 0),   // stage 3: 8 วัน
                   TimeSpan(0, 0, 0, 10),  // LED on: 12 ชั่วโมง
                   TimeSpan(0, 0, 0, 5),  // LED off: 12 ชั่วโมง
                   15, 2, 4, "Tomato", "Floor 1");

GrowSystem system2(TimeSpan(0, 0, 0, 30),  // total: 45 วัน
                   TimeSpan(0, 0, 0, 10),   // stage 1: 7 วัน
                   TimeSpan(0, 0, 0, 10),  // stage 2: 30 วัน
                   TimeSpan(0, 0, 0, 10),   // stage 3: 8 วัน
                   TimeSpan(0, 0, 0, 2),  // LED on: 12 ชั่วโมง
                   TimeSpan(0, 0, 0, 4),  // LED off: 12 ชั่วโมง
                   13, 12, 14, "Egg", "Floor 2");

void IRAM_ATTR ISR_button1() {
    if ((millis() - lastDebounceTime1) > DEBOUNCE_DELAY) {
        button1Pressed = true;
        lastDebounceTime1 = millis();
    }
}

void IRAM_ATTR ISR_button2() {
    if ((millis() - lastDebounceTime2) > DEBOUNCE_DELAY) {
        button2Pressed = true;
        lastDebounceTime2 = millis();
    }
}

void setup() {
    Serial.begin(115200);
    
    pinMode(system1.startButtonPin, INPUT_PULLUP);
    pinMode(system1.ledPin, OUTPUT);
    pinMode(system1.fanPin, OUTPUT);
    pinMode(system2.startButtonPin, INPUT_PULLUP);
    pinMode(system2.ledPin, OUTPUT);
    pinMode(system2.fanPin, OUTPUT);
    
    attachInterrupt(digitalPinToInterrupt(system1.startButtonPin), ISR_button1, FALLING);
    attachInterrupt(digitalPinToInterrupt(system2.startButtonPin), ISR_button2, FALLING);

    if (!rtc.begin()) {
        Serial.println("Couldn't find RTC");
        while (1);
    }

    if (rtc.lostPower()) {
        Serial.println("RTC lost power, lets set the time!");
        rtc.adjust(DateTime(F(__DATE__), F(__TIME__)));
    }
    
    if(!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
        Serial.println(F("SSD1306 allocation failed"));
        for(;;);
    }
    
    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.display();
}

void loop() {
    handleButtons();
    updateSystems();
    updateDisplay();
}

void handleButtons() {
    if (button1Pressed) {
        button1Pressed = false;
        if (!system1.countdownStarted && !system1.countdownOver) {
            system1.start();
        } else if (system1.countdownOver) {
            system1.reset();
        }
    }
    
    if (button2Pressed) {
        button2Pressed = false;
        if (!system2.countdownStarted && !system2.countdownOver) {
            system2.start();
        } else if (system2.countdownOver) {
            system2.reset();
        }
    }
}

void updateSystems() {
    system1.update();
    system2.update();
}

void updateDisplay() {
    display.clearDisplay();
    system1.updateDisplay(0);
    system2.updateDisplay(32);
    display.display();
}