#define LOAD_COUNT 4

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

const char* ssid = "Wokwi-GUEST";
const char* password = "";

String DEVICE_NAME = "4_module_try";
const byte loads[LOAD_COUNT] = {19,21,22,23};
const byte switches[LOAD_COUNT] = {13, 14, 12, 27};
const byte Push_button = 36;

const char* mqtt_server = "broker.emqx.io";
const char* dtopic = "module001/fib001/dev001";
const char* ptopic = "module001/fib001/sta001";
const char* stopic = "module001/fib001/lo001";
const char* willMessage = "0";
const boolean retained = true;

boolean loadStates[LOAD_COUNT] = {HIGH, HIGH, HIGH, HIGH}; // Default states for loads
boolean switchStates[LOAD_COUNT] = {HIGH, HIGH, HIGH, HIGH};

WiFiClient wifiClient;
PubSubClient mqttClient(wifiClient);

String previous_result;
String present_result;

void setup_wifi() {
    Serial.println("Connecting to WiFi...");
    WiFi.begin(ssid, password);

    while (WiFi.status() != WL_CONNECTED) {
        delay(500);
        Serial.print(".");
    }
    Serial.println("\nWiFi connected");
    Serial.print("IP address: ");
    Serial.println(WiFi.localIP());
}

void mqtt_callback(char* topic, byte* payload, unsigned int length) {
    String receivedPayload((char*)payload, length);
    Serial.printf("Message received on topic %s: %s\n", topic, receivedPayload.c_str());

    StaticJsonDocument<200> doc;
    DeserializationError error = deserializeJson(doc, receivedPayload);

    if (error) {
        Serial.print(F("JSON parsing failed: "));
        Serial.println(error.f_str());
        return;
    }

    const char* state = doc["state"];
    int loadNo = doc["load"];

    if (loadNo >= 1 && loadNo <= LOAD_COUNT) {
        int loadIdx = loadNo - 1;
        if (strcmp(state, "on") == 0) {
            digitalWrite(loads[loadIdx], LOW);
            loadStates[loadIdx] = LOW;
            Serial.printf("Load %d turned ON\n", loadNo);
        } else if (strcmp(state, "off") == 0) {
            digitalWrite(loads[loadIdx], HIGH);
            loadStates[loadIdx] = HIGH;
            Serial.printf("Load %d turned OFF\n", loadNo);
        } else {
            Serial.printf("Invalid state for load %d\n", loadNo);
        }
    } else {
        Serial.println("Invalid load number received in MQTT payload");
    }
}

void reconnect_mqtt() {
    while (!mqttClient.connected()) {
        Serial.print("Attempting MQTT connection...");
        String clientId = DEVICE_NAME + String(random(0xffff), HEX);
        if (mqttClient.connect(clientId.c_str(), dtopic, 0, retained, String(willMessage).c_str())) {
            Serial.println("connected");
            const char* message = "1";
            int length = strlen(message);
            mqttClient.publish(dtopic, (byte*)message, length, retained);
            mqttClient.subscribe(stopic);
        } else {
            Serial.print("failed, rc=");
            Serial.print(mqttClient.state());
            Serial.println(" try again in 5 seconds");
            delay(5000);
        }
    }
}

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

    for (int i = 0; i < LOAD_COUNT; i++) {
        pinMode(loads[i], OUTPUT);
        pinMode(switches[i], INPUT_PULLUP);
        digitalWrite(loads[i], HIGH); // Turn loads OFF initially
    }
    pinMode(Push_button, INPUT_PULLUP);

    setup_wifi();

    mqttClient.setServer(mqtt_server, 1883);
    mqttClient.setCallback(mqtt_callback);
}

void handleLocalSwitches() {
    for (int i = 0; i < LOAD_COUNT; i++) {
        boolean currentSwitchState = digitalRead(switches[i]);
        if (currentSwitchState != switchStates[i]) {
            // State change detected
            loadStates[i] = !loadStates[i];
            digitalWrite(loads[i], loadStates[i] ? HIGH : LOW);

            // Publish state change to MQTT
            StaticJsonDocument<200> doc;
            doc["load"] = i + 1;
            doc["state"] = loadStates[i] ? "off" : "on";

            String payload;
            serializeJson(doc, payload);
            mqttClient.publish(ptopic, payload.c_str(), retained);

            Serial.printf("Switch %d toggled. Load %d is now %s\n", i + 1, i + 1, loadStates[i] ? "OFF" : "ON");
            switchStates[i] = currentSwitchState;
        }
    }
}

void publishLoadStates() {
    StaticJsonDocument<200> doc;
    for (int i = 0; i < LOAD_COUNT; i++) {
        doc["Load" + String(i + 1)] = loadStates[i] ? "off" : "on";
    }

    String payload;
    serializeJson(doc, payload);
    if (payload != previous_result) {
        mqttClient.publish(ptopic, payload.c_str(), retained);
        previous_result = payload;
    }
}

void loop() {
    if (WiFi.status() != WL_CONNECTED) {
        setup_wifi();
    }

    if (!mqttClient.connected()) {
        reconnect_mqtt();
    }

    handleLocalSwitches();
    publishLoadStates();
    mqttClient.loop();
}