#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <avr/interrupt.h>

// OLED configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define OLED_ADDRESS 0x3C

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);
bool oled_ok = false;

// Pin definitions
#define SOIL_PIN A0
#define TEMP_PIN A1

#define BUTTON_PIN 2
#define RELAY_PIN 4

#define LED_RED   5
#define LED_GREEN 6
#define LED_BLUE  7

// Soil calibration
#define SOIL_DRY_ADC 400
#define SOIL_WET_ADC 250
#define SOIL_DRY_ADC_START 400
#define SOIL_WET_ADC_STOP 300
#define SOIL_STOP_MOISTURE_PERCENT 70

#define STARTUP_STABILIZE_MS 10000
#define DRY_CONFIRM_READINGS 4
#define WET_CONFIRM_READINGS 3
#define MIN_WATERING_BEFORE_EARLY_STOP_MS 1500

#define BUTTON_DEBOUNCE_MS 300

#define RELAY_ON  LOW
#define RELAY_OFF HIGH

#define PUMP_FLOW_ML_PER_SEC 5.0
#define MAX_WATERING_TIME_MS 8000
#define WAIT_AFTER_WATERING_MS 30000
#define MANUAL_WATERING_TIME_MS 4000

#define STATUS_PRINT_INTERVAL_MS 3000

enum SystemState {
    SYSTEM_STABILIZING,
    SYSTEM_OK,
    NEEDS_WATER,
    WATERING,
    WAIT_AFTER_WATERING
};

SystemState state = SYSTEM_STABILIZING;

unsigned long watering_start_ms = 0;
unsigned long watering_duration_ms = 0;
unsigned long wait_start_ms = 0;
unsigned long system_start_ms = 0;
unsigned long last_status_print_ms = 0;
unsigned long last_button_press_ms = 0;

int wet_reading_count = 0;
int dry_reading_count = 0;
bool manual_watering_active = false;

// Interrupt flag
volatile bool button_interrupt_flag = false;

ISR(INT0_vect)
{
    button_interrupt_flag = true;
}

void initButtonInterrupt()
{
    pinMode(BUTTON_PIN, INPUT_PULLUP);

    cli();

    // INT0 on falling edge:
    // button released = HIGH, button pressed = LOW
    EICRA |= (1 << ISC01);
    EICRA &= ~(1 << ISC00);

    // Enable INT0
    EIMSK |= (1 << INT0);

    sei();
}

const char* stateToString(SystemState current_state)
{
    switch (current_state)
    {
        case SYSTEM_STABILIZING: return "STABILIZING";
        case SYSTEM_OK: return "OK";
        case NEEDS_WATER: return "NEEDS WATER";
        case WATERING:
            if (manual_watering_active)
                return "MANUAL WATERING";
            return "AUTO WATERING";
        case WAIT_AFTER_WATERING: return "WAITING FOR WATER";
        default: return "UNKNOWN";
    }
}

const char* soilConditionToString(int soil_adc)
{
    if (soil_adc >= SOIL_DRY_ADC_START)
        return "DRY";

    if (soil_adc <= SOIL_WET_ADC_STOP)
        return "WET";

    return "NORMAL";
}

void ledsOff()
{
    analogWrite(LED_RED, 0);
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_BLUE, LOW);
}

void showOkLed()
{
    analogWrite(LED_RED, 0);
    digitalWrite(LED_BLUE, LOW);
    digitalWrite(LED_GREEN, HIGH);
}

void showManualWateringLed()
{
    analogWrite(LED_RED, 0);
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_BLUE, HIGH);
}

void showAutomaticWateringLed(unsigned long now)
{
    unsigned long elapsed = now - watering_start_ms;

    if (elapsed > watering_duration_ms)
        elapsed = watering_duration_ms;

    int red_brightness = map(elapsed, 0, watering_duration_ms, 255, 0);
    red_brightness = constrain(red_brightness, 0, 255);

    analogWrite(LED_RED, red_brightness);
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_BLUE, HIGH);
}

void showNeedWaterLed()
{
    analogWrite(LED_RED, 255);
    digitalWrite(LED_GREEN, LOW);
    digitalWrite(LED_BLUE, LOW);
}

void updateLedOutputs(unsigned long now)
{
    if (state == SYSTEM_STABILIZING)
    {
        ledsOff();
    }
    else if (state == WATERING)
    {
        if (manual_watering_active)
            showManualWateringLed();
        else
            showAutomaticWateringLed(now);
    }
    else if (state == WAIT_AFTER_WATERING)
    {
        showOkLed();
    }
    else if (state == NEEDS_WATER)
    {
        showNeedWaterLed();
    }
    else
    {
        showOkLed();
    }
}

void pumpOn()
{
    digitalWrite(RELAY_PIN, RELAY_ON);
}

void pumpOff()
{
    digitalWrite(RELAY_PIN, RELAY_OFF);
}

int readSoilAdcAverage()
{
    long sum = 0;
    const int samples = 10;

    for (int i = 0; i < samples; i++)
    {
        sum += analogRead(SOIL_PIN);
        delay(5);
    }

    return sum / samples;
}

int readMoisturePercent(int soil_adc)
{
    int moisture_percent = map(soil_adc, SOIL_DRY_ADC, SOIL_WET_ADC, 0, 100);
    moisture_percent = constrain(moisture_percent, 0, 100);

    return moisture_percent;
}

float readTemperatureC(int temp_adc)
{
    float voltage = temp_adc * 3.3 / 1024.0;
    return voltage * 100.0;
}

float calculateWaterAmountMl(int moisture_percent, float temp_c)
{
    float water_ml = 0.0;

    if (moisture_percent < 15)
        water_ml = 25.0;
    else if (moisture_percent < 30)
        water_ml = 18.0;
    else if (moisture_percent < 40)
        water_ml = 10.0;

    if (water_ml > 0)
    {
        if (temp_c > 28.0)
            water_ml += 5.0;
        else if (temp_c < 20.0)
            water_ml -= 3.0;
    }

    if (water_ml < 0)
        water_ml = 0;

    if (water_ml > 30)
        water_ml = 30;

    return water_ml;
}

unsigned long calculatePumpTimeMs(float water_ml)
{
    if (water_ml <= 0)
        return 0;

    float time_sec = water_ml / PUMP_FLOW_ML_PER_SEC;
    unsigned long time_ms = (unsigned long)(time_sec * 1000.0);

    if (time_ms > MAX_WATERING_TIME_MS)
        time_ms = MAX_WATERING_TIME_MS;

    return time_ms;
}

void initOLED()
{
    oled_ok = display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDRESS);

    if (!oled_ok)
    {
        Serial.println(F("[SYSTEM] Display not detected."));
        return;
    }

    display.clearDisplay();
    display.setTextSize(1);
    display.setTextColor(SSD1306_WHITE);
    display.setCursor(0, 0);
    display.println(F("Smart Orchid"));
    display.setCursor(0, 16);
    display.println(F("Starting system"));
    display.display();

    Serial.println(F("[SYSTEM] Display initialized."));
}

void updateOLED(int soil_adc, int moisture_percent, float temp_c)
{
    if (!oled_ok)
        return;

    display.clearDisplay();
    display.setTextColor(SSD1306_WHITE);
    display.setTextSize(1);

    display.setCursor(0, 0);
    display.println(F("Smart Orchid"));
    display.drawLine(0, 10, 127, 10, SSD1306_WHITE);

    display.setCursor(0, 14);

    if (state == SYSTEM_STABILIZING)
        display.println(F("Status: STARTING"));
    else if (state == SYSTEM_OK)
        display.println(F("Status: HEALTHY"));
    else if (state == NEEDS_WATER)
        display.println(F("Status: DRY SOIL"));
    else if (state == WATERING && manual_watering_active)
        display.println(F("Status: MANUAL WATER"));
    else if (state == WATERING)
        display.println(F("Status: WATERING"));
    else if (state == WAIT_AFTER_WATERING)
        display.println(F("Status: ABSORBING"));

    display.setCursor(0, 28);
    display.print(F("Moist: "));
    display.print(moisture_percent);
    display.print(F("%"));

    display.setCursor(70, 28);
    display.print(F("ADC:"));
    display.print(soil_adc);

    display.setCursor(0, 40);
    display.print(F("Temp: "));
    display.print(temp_c, 1);
    display.println(F(" C"));

    display.setCursor(0, 52);

    if (state == SYSTEM_STABILIZING)
    {
        unsigned long elapsed = millis() - system_start_ms;
        unsigned long remaining = 0;

        if (elapsed < STARTUP_STABILIZE_MS)
            remaining = (STARTUP_STABILIZE_MS - elapsed) / 1000;

        display.print(F("Ready in "));
        display.print(remaining);
        display.print(F("s"));
    }
    else if (state == SYSTEM_OK)
    {
        display.print(F("Soil: "));
        display.print(soilConditionToString(soil_adc));
    }
    else if (state == NEEDS_WATER)
    {
        display.print(F("Preparing pump"));
    }
    else if (state == WATERING)
    {
        display.print(F("Pump running"));
    }
    else if (state == WAIT_AFTER_WATERING)
    {
        display.print(F("Water absorbing"));
    }

    display.display();
}

void startWatering(unsigned long duration_ms, bool manual)
{
    if (duration_ms == 0)
        return;

    state = WATERING;
    watering_start_ms = millis();
    watering_duration_ms = duration_ms;

    wet_reading_count = 0;
    dry_reading_count = 0;
    manual_watering_active = manual;

    pumpOn();

    if (manual_watering_active)
        Serial.print(F("[SYSTEM] Manual watering started for "));
    else
        Serial.print(F("[SYSTEM] Automatic watering started for "));

    Serial.print(watering_duration_ms);
    Serial.println(F(" ms."));
}

void stopWatering()
{
    pumpOff();

    manual_watering_active = false;
    wet_reading_count = 0;
    dry_reading_count = 0;

    Serial.println(F("[SYSTEM] Watering stopped."));
}

void printSystemStatus(int soil_adc, int moisture_percent, float temp_c)
{
    Serial.print(F("[SYSTEM] State: "));
    Serial.print(stateToString(state));

    Serial.print(F(" | Soil ADC: "));
    Serial.print(soil_adc);

    Serial.print(F(" | Soil: "));
    Serial.print(soilConditionToString(soil_adc));

    Serial.print(F(" | Moisture: "));
    Serial.print(moisture_percent);
    Serial.print(F("%"));

    Serial.print(F(" | Temperature: "));
    Serial.print(temp_c, 1);
    Serial.print(F(" C"));

    Serial.println();
}

void handleManualWateringButton(unsigned long now)
{
    bool button_event = false;

    cli();
    if (button_interrupt_flag)
    {
        button_interrupt_flag = false;
        button_event = true;
    }
    sei();

    if (!button_event)
        return;

    if (now - last_button_press_ms < BUTTON_DEBOUNCE_MS)
        return;

    last_button_press_ms = now;

    if (state == WATERING)
        return;

    Serial.println(F("[SYSTEM] Manual watering requested."));
    startWatering(MANUAL_WATERING_TIME_MS, true);
}

void setup()
{
    Serial.begin(9600);

    pinMode(RELAY_PIN, OUTPUT);
    digitalWrite(RELAY_PIN, RELAY_OFF);

    pinMode(LED_RED, OUTPUT);
    pinMode(LED_GREEN, OUTPUT);
    pinMode(LED_BLUE, OUTPUT);

    ledsOff();
    pumpOff();

    initButtonInterrupt();

    delay(1000);

    Serial.println(F("Smart Orchid self-watering system started."));
    Serial.println(F("Sensors are stabilizing. Pump and LEDs are off."));
    Serial.println(F("Soil calibration: dry >= 400 ADC, wet <= 300 ADC."));
    Serial.println(F("Button interrupt enabled on INT0."));

    initOLED();

    system_start_ms = millis();
    state = SYSTEM_STABILIZING;
}

void loop()
{
    unsigned long now = millis();

    handleManualWateringButton(now);

    int soil_adc = readSoilAdcAverage();
    int temp_adc = analogRead(TEMP_PIN);

    int moisture_percent = readMoisturePercent(soil_adc);
    float temp_c = readTemperatureC(temp_adc);

    if (state == SYSTEM_STABILIZING)
    {
        pumpOff();
        ledsOff();

        if (now - last_status_print_ms >= STATUS_PRINT_INTERVAL_MS)
        {
            printSystemStatus(soil_adc, moisture_percent, temp_c);
            last_status_print_ms = now;
        }

        updateOLED(soil_adc, moisture_percent, temp_c);

        if (now - system_start_ms >= STARTUP_STABILIZE_MS)
        {
            state = SYSTEM_OK;
            dry_reading_count = 0;
            wet_reading_count = 0;

            Serial.println(F("[SYSTEM] Stabilization finished. Monitoring started."));
        }

        delay(500);
        return;
    }

    if (state == WATERING)
    {
        unsigned long elapsed_watering = now - watering_start_ms;

        bool early_stop_allowed =
            elapsed_watering >= MIN_WATERING_BEFORE_EARLY_STOP_MS;

        bool soil_is_wet_enough =
            (soil_adc <= SOIL_WET_ADC_STOP) ||
            (moisture_percent >= SOIL_STOP_MOISTURE_PERCENT);

        if (!manual_watering_active && early_stop_allowed && soil_is_wet_enough)
        {
            wet_reading_count++;
        }
        else
        {
            wet_reading_count = 0;
        }

        if (!manual_watering_active && wet_reading_count >= WET_CONFIRM_READINGS)
        {
            Serial.println(F("[SYSTEM] Soil moisture is sufficient. Automatic watering will stop."));

            stopWatering();

            state = WAIT_AFTER_WATERING;
            wait_start_ms = now;

            Serial.println(F("[SYSTEM] Waiting for water to absorb into the soil."));
        }
        else if (now - watering_start_ms >= watering_duration_ms)
        {
            stopWatering();

            state = WAIT_AFTER_WATERING;
            wait_start_ms = now;

            Serial.println(F("[SYSTEM] Waiting for water to absorb into the soil."));
        }
    }
    else if (state == WAIT_AFTER_WATERING)
    {
        if (now - wait_start_ms >= WAIT_AFTER_WATERING_MS)
        {
            state = SYSTEM_OK;
            dry_reading_count = 0;
            wet_reading_count = 0;

            Serial.println(F("[SYSTEM] Absorption period finished. Monitoring resumed."));
        }
    }
    else
    {
        bool soil_is_dry_for_start = (soil_adc >= SOIL_DRY_ADC_START);

        if (soil_is_dry_for_start)
        {
            dry_reading_count++;
            state = SYSTEM_OK;

            if (dry_reading_count >= DRY_CONFIRM_READINGS)
            {
                Serial.println(F("[SYSTEM] Soil is consistently dry. Automatic watering will start."));

                state = NEEDS_WATER;
                dry_reading_count = 0;

                float water_ml = calculateWaterAmountMl(moisture_percent, temp_c);
                unsigned long pump_time_ms = calculatePumpTimeMs(water_ml);

                if (pump_time_ms > 0)
                    startWatering(pump_time_ms, false);
            }
        }
        else
        {
            dry_reading_count = 0;
            state = SYSTEM_OK;
        }
    }

    if (now - last_status_print_ms >= STATUS_PRINT_INTERVAL_MS)
    {
        printSystemStatus(soil_adc, moisture_percent, temp_c);
        last_status_print_ms = now;
    }

    updateLedOutputs(now);
    updateOLED(soil_adc, moisture_percent, temp_c);

    delay(500);
}