#include <WiFi.h>
#include <PubSubClient.h>
#include <DHTesp.h>

const char *WIFI_SSID = "Wokwi-GUEST";
const char *WIFI_PASSWORD = "";

const char *HA_BROKER = "broker.mqttdashboard.com";
const uint16_t MQTT_PORT = 1883;
const char *MQTT_CLIENT_ID = "hyhung-kitchen-wokwi";

const char *STATUS_TOPIC = "hyhung/kitchen/status";
const char *TEMPERATURE_TOPIC = "hyhung/kitchen/temperature";
const char *HUMIDITY_TOPIC = "hyhung/kitchen/humidity";
const char *BATTERY_TOPIC = "hyhung/kitchen/battery_soc";
const char *HOUSE_POWER_TOPIC = "hyhung/kitchen/house_power";

const uint8_t DHT_PIN = 4;

const uint8_t applianceButtonPins[] = {13, 12, 14};
const uint8_t applianceLedPins[] = {5, 18, 19};
const char *applianceNames[] = {"dryer", "washing_machine", "dishwasher"};
const char *applianceStateTopics[] = {
    "hyhung/kitchen/dryer/state",
    "hyhung/kitchen/washing_machine/state",
    "hyhung/kitchen/dishwasher/state"};
const char *applianceCommandTopics[] = {
    "hyhung/kitchen/dryer/command",
    "hyhung/kitchen/washing_machine/command",
    "hyhung/kitchen/dishwasher/command"};
const char *appliancePowerTopics[] = {
    "hyhung/kitchen/dryer/power",
    "hyhung/kitchen/washing_machine/power",
    "hyhung/kitchen/dishwasher/power"};
const char *applianceRemainingTopics[] = {
    "hyhung/kitchen/dryer/remaining",
    "hyhung/kitchen/washing_machine/remaining",
    "hyhung/kitchen/dishwasher/remaining"};
const int appliancePowerWatts[] = {1800, 1200, 1500};
const unsigned long applianceRunMs[] = {3600000UL, 1800000UL, 3600000UL};

const uint8_t lightButtonPins[] = {27, 26, 25};
const uint8_t relayPins[] = {21, 22, 23};
const char *lightNames[] = {"poppy_light", "cooker_light", "sink_light"};
const char *lightStateTopics[] = {
    "hyhung/kitchen/poppy_light/state",
    "hyhung/kitchen/cooker_light/state",
    "hyhung/kitchen/sink_light/state"};
const char *lightCommandTopics[] = {
    "hyhung/kitchen/poppy_light/command",
    "hyhung/kitchen/cooker_light/command",
    "hyhung/kitchen/sink_light/command"};
const char *lightPowerTopics[] = {
    "hyhung/kitchen/poppy_light/power",
    "hyhung/kitchen/cooker_light/power",
    "hyhung/kitchen/sink_light/power"};
const int lightPowerWatts[] = {5, 10, 8};

const size_t applianceCount = sizeof(applianceButtonPins) / sizeof(applianceButtonPins[0]);
const size_t lightCount = sizeof(lightButtonPins) / sizeof(lightButtonPins[0]);

bool applianceStates[3] = {false, false, false};
bool lightStates[3] = {false, false, false};
bool lastApplianceButtonRead[3] = {true, true, true};
bool lastLightButtonRead[3] = {true, true, true};
unsigned long applianceStartMs[3] = {0, 0, 0};
unsigned long lastButtonDebounceMs[6] = {0, 0, 0, 0, 0, 0};

WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);
DHTesp dhtSensor;

unsigned long lastSensorPublishMs = 0;
unsigned long lastHeartbeatMs = 0;

void logLine(const String &message)
{
    Serial.println(message);
}

void logTopicValue(const char *prefix, const char *topic, const String &value)
{
    Serial.print(prefix);
    Serial.print(" ");
    Serial.print(topic);
    Serial.print(" = ");
    Serial.println(value);
}

void publishText(const char *topic, const char *payload, bool retained = true)
{
    logTopicValue("MQTT publish", topic, payload);
    mqttClient.publish(topic, payload, retained);
}

void publishNumber(const char *topic, float value, bool retained = true)
{
    char payload[32];
    dtostrf(value, 0, 2, payload);
    logTopicValue("MQTT publish", topic, payload);
    mqttClient.publish(topic, payload, retained);
}

void setApplianceState(size_t index, bool active)
{
    applianceStates[index] = active;
    digitalWrite(applianceLedPins[index], active ? HIGH : LOW);
    applianceStartMs[index] = active ? millis() : 0;
    publishText(applianceStateTopics[index], active ? "ON" : "OFF");
    publishNumber(appliancePowerTopics[index], active ? appliancePowerWatts[index] : 0);
    publishNumber(applianceRemainingTopics[index], active ? (applianceRunMs[index] / 1000UL) : 0);
}

void setLightState(size_t index, bool active)
{
    lightStates[index] = active;
    digitalWrite(relayPins[index], active ? HIGH : LOW);
    publishText(lightStateTopics[index], active ? "ON" : "OFF");
    publishNumber(lightPowerTopics[index], active ? lightPowerWatts[index] : 0);
}

void publishStateSet()
{
    for (size_t i = 0; i < applianceCount; i++)
    {
        publishText(applianceStateTopics[i], applianceStates[i] ? "ON" : "OFF");
        publishNumber(appliancePowerTopics[i], applianceStates[i] ? appliancePowerWatts[i] : 0);
        unsigned long remaining = 0;
        if (applianceStates[i] && applianceStartMs[i] > 0)
        {
            unsigned long elapsed = millis() - applianceStartMs[i];
            if (elapsed < applianceRunMs[i])
            {
                remaining = (applianceRunMs[i] - elapsed) / 1000UL;
            }
        }
        publishNumber(applianceRemainingTopics[i], remaining);
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        publishText(lightStateTopics[i], lightStates[i] ? "ON" : "OFF");
        publishNumber(lightPowerTopics[i], lightStates[i] ? lightPowerWatts[i] : 0);
    }
}

void handleCommand(const char *topic, const char *payload)
{
    logTopicValue("MQTT rx", topic, payload);
    String command = String(payload);
    command.trim();
    command.toUpperCase();

    for (size_t i = 0; i < applianceCount; i++)
    {
        if (strcmp(topic, applianceCommandTopics[i]) == 0)
        {
            if (command == "ON")
            {
                setApplianceState(i, true);
            }
            else if (command == "OFF")
            {
                setApplianceState(i, false);
            }
            else if (command == "TOGGLE")
            {
                setApplianceState(i, !applianceStates[i]);
            }
            return;
        }
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        if (strcmp(topic, lightCommandTopics[i]) == 0)
        {
            if (command == "ON")
            {
                setLightState(i, true);
            }
            else if (command == "OFF")
            {
                setLightState(i, false);
            }
            else if (command == "TOGGLE")
            {
                setLightState(i, !lightStates[i]);
            }
            return;
        }
    }
}

void mqttCallback(char *topic, byte *payload, unsigned int length)
{
    char message[64];
    unsigned int copyLength = length < sizeof(message) - 1 ? length : sizeof(message) - 1;
    memcpy(message, payload, copyLength);
    message[copyLength] = '\0';
    handleCommand(topic, message);
}

void connectWiFi()
{
    logLine("WiFi: connecting to " + String(WIFI_SSID));
    WiFi.mode(WIFI_STA);
    WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

    while (WiFi.status() != WL_CONNECTED)
    {
        delay(500);
    }

    logLine("WiFi: connected, IP = " + WiFi.localIP().toString());
}

void subscribeTopics()
{
    for (size_t i = 0; i < applianceCount; i++)
    {
        mqttClient.subscribe(applianceCommandTopics[i]);
        logLine("MQTT subscribe " + String(applianceCommandTopics[i]));
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        mqttClient.subscribe(lightCommandTopics[i]);
        logLine("MQTT subscribe " + String(lightCommandTopics[i]));
    }
}

void reconnectMqtt()
{
    logLine("MQTT: connecting to " + String(HA_BROKER) + ":" + String(MQTT_PORT));
    while (!mqttClient.connected())
    {
        if (mqttClient.connect(MQTT_CLIENT_ID, nullptr, nullptr, STATUS_TOPIC, 0, true, "offline"))
        {
            logLine("MQTT: connected as " + String(MQTT_CLIENT_ID));
            publishText(STATUS_TOPIC, "online");
            subscribeTopics();
            publishStateSet();
        }
        else
        {
            Serial.print("MQTT connect failed, rc=");
            Serial.print(mqttClient.state());
            Serial.println(" retrying...");
            delay(1500);
        }
    }
}

void publishEnvironment()
{
    TempAndHumidity reading = dhtSensor.getTempAndHumidity();
    float temperature = reading.temperature;
    float humidity = reading.humidity;

    if (isnan(temperature))
    {
        temperature = 23.5 + 1.3 * sin(millis() / 60000.0);
    }

    if (isnan(humidity))
    {
        humidity = 48.0 + 5.0 * sin(millis() / 45000.0 + 1.2);
    }

    float batterySoc = 76.0 + 12.0 * sin(millis() / 120000.0 + 0.5);
    float housePower = 120.0;

    for (size_t i = 0; i < applianceCount; i++)
    {
        if (applianceStates[i])
        {
            housePower += appliancePowerWatts[i];
        }
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        if (lightStates[i])
        {
            housePower += lightPowerWatts[i];
        }
    }

    publishNumber(TEMPERATURE_TOPIC, temperature);
    publishNumber(HUMIDITY_TOPIC, humidity);
    publishNumber(BATTERY_TOPIC, batterySoc);
    publishNumber(HOUSE_POWER_TOPIC, housePower);

    Serial.print("Sensor snapshot -> T:");
    Serial.print(temperature, 2);
    Serial.print(" C, H:");
    Serial.print(humidity, 2);
    Serial.print(" %, Battery:");
    Serial.print(batterySoc, 2);
    Serial.print(" %, Power:");
    Serial.print(housePower, 2);
    Serial.println(" W");
}

void pollButtons()
{
    const unsigned long debounceMs = 35;

    for (size_t i = 0; i < applianceCount; i++)
    {
        bool currentRead = digitalRead(applianceButtonPins[i]);
        if (currentRead != lastApplianceButtonRead[i])
        {
            lastButtonDebounceMs[i] = millis();
            lastApplianceButtonRead[i] = currentRead;
        }

        if ((millis() - lastButtonDebounceMs[i]) > debounceMs && currentRead == LOW)
        {
            while (digitalRead(applianceButtonPins[i]) == LOW)
            {
                delay(10);
            }
            setApplianceState(i, !applianceStates[i]);
        }
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        size_t index = i + applianceCount;
        bool currentRead = digitalRead(lightButtonPins[i]);
        if (currentRead != lastLightButtonRead[i])
        {
            lastButtonDebounceMs[index] = millis();
            lastLightButtonRead[i] = currentRead;
        }

        if ((millis() - lastButtonDebounceMs[index]) > debounceMs && currentRead == LOW)
        {
            while (digitalRead(lightButtonPins[i]) == LOW)
            {
                delay(10);
            }
            setLightState(i, !lightStates[i]);
        }
    }
}

void updateApplianceCountdowns()
{
    for (size_t i = 0; i < applianceCount; i++)
    {
        if (!applianceStates[i] || applianceStartMs[i] == 0)
        {
            continue;
        }

        unsigned long elapsed = millis() - applianceStartMs[i];
        if (elapsed >= applianceRunMs[i])
        {
            setApplianceState(i, false);
            continue;
        }

        unsigned long remainingSeconds = (applianceRunMs[i] - elapsed) / 1000UL;
        publishNumber(applianceRemainingTopics[i], remainingSeconds);
    }
}

void setup()
{
    Serial.begin(115200);
    delay(200);

    logLine("Kitchen Dashboard firmware booting...");

    for (size_t i = 0; i < applianceCount; i++)
    {
        pinMode(applianceButtonPins[i], INPUT_PULLUP);
        pinMode(applianceLedPins[i], OUTPUT);
        digitalWrite(applianceLedPins[i], LOW);
    }

    for (size_t i = 0; i < lightCount; i++)
    {
        pinMode(lightButtonPins[i], INPUT_PULLUP);
        pinMode(relayPins[i], OUTPUT);
        digitalWrite(relayPins[i], LOW);
    }

    dhtSensor.setup(DHT_PIN, DHTesp::DHT22);

    connectWiFi();

    mqttClient.setServer(HA_BROKER, MQTT_PORT);
    mqttClient.setCallback(mqttCallback);
    logLine("Setup complete");
}

void loop()
{
    if (!mqttClient.connected())
    {
        reconnectMqtt();
    }

    mqttClient.loop();
    pollButtons();
    updateApplianceCountdowns();

    if (millis() - lastSensorPublishMs > 5000)
    {
        lastSensorPublishMs = millis();
        publishEnvironment();
        publishStateSet();
    }

    if (millis() - lastHeartbeatMs > 10000)
    {
        lastHeartbeatMs = millis();
        Serial.print("Heartbeat -> WiFi:");
        Serial.print(WiFi.status() == WL_CONNECTED ? "connected" : "down");
        Serial.print(", MQTT:");
        Serial.print(mqttClient.connected() ? "connected" : "down");
        Serial.print(", uptime_s:");
        Serial.println(millis() / 1000UL);
    }
}