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

// ===== Wi-Fi Configuration =====
const char* ssid     = "Wokwi-GUEST";
const char* password = "";

// ===== MQTT Configuration =====
const char* mqttHost     = "public.cloud.shiftr.io";
const int   mqttPort     = 1883;  // Non-TLS
const char* mqttUser     = "public";
const char* mqttPass     = "public";
const char* mqttClientID = "esp32Client";
const char* topic        = "chat/room1";

// ===== Hardware Configuration =====
struct LED {
  int pin;
  String name;
};

LED leds[] = {
  {2, "led1"},
  {4, "led2"},
  {22, "led3"},
  {23, "led4"}
};

const int ledCount = sizeof(leds) / sizeof(leds[0]);

// ===== Global Variables =====
WiFiClient wifiClient;
PubSubClient mqtt(wifiClient);
unsigned long lastPing = 0;
const unsigned long pingInterval = 1500; // milliseconds

// ===== Helper Functions =====
void setLED(const String &cmd, LED &led, bool mqttCmd = true) {
  if (cmd.equalsIgnoreCase(led.name + " on")) {
    digitalWrite(led.pin, HIGH);
    Serial.println(led.name + " turned ON" + (mqttCmd ? " via MQTT" : " locally"));
  } 
  else if (cmd.equalsIgnoreCase(led.name + " off")) {
    digitalWrite(led.pin, LOW);
    Serial.println(led.name + " turned OFF" + (mqttCmd ? " via MQTT" : " locally"));
  }
}

String processMessage(String msg) {
    const char* leds[] = {"led1", "led2", "led3", "led4"};
    const char* states[] = {"on", "off"};
    int numLeds = sizeof(leds) / sizeof(leds[0]);
    int numStates = sizeof(states) / sizeof(states[0]);

    for (int i = 0; i < numLeds; i++) {
        for (int j = 0; j < numStates; j++) {
            String command = "!" + String(leds[i]) + " " + String(states[j]);
            if (msg.indexOf(command) >= 0) {
                msg = String(leds[i]) + " " + String(states[j]);
                return msg; // Return immediately after first match
            }
        }
    }

    return msg; // No match, return original message
}


// ===== MQTT Callback =====
void callback(char* topic, byte* payload, unsigned int length) {
  String msg;
  for (unsigned int i = 0; i < length; i++) msg += (char)payload[i];
  msg.trim();

  if (msg.equalsIgnoreCase("~ping~") == false) {
    Serial.println("Received: " + msg);
  }

  msg=processMessage(msg);

  // Control LEDs
  for (int i = 0; i < ledCount; i++) {
    setLED(msg, leds[i]);
  }
}

// ===== MQTT Connection =====
void connectMQTT() {
  while (!mqtt.connected()) {
    Serial.print("Connecting MQTT...");
    if (mqtt.connect(mqttClientID, mqttUser, mqttPass)) {
      Serial.println("connected!");
      mqtt.subscribe(topic);
      Serial.println("Subscribed to " + String(topic));
    } else {
      Serial.print("failed, rc=");
      Serial.println(mqtt.state());
      delay(2000);
    }
  }
}

// ===== Setup =====
void setup() {
  Serial.begin(115200);

  // Initialize LEDs
  for (int i = 0; i < ledCount; i++) {
    pinMode(leds[i].pin, OUTPUT);
    digitalWrite(leds[i].pin, LOW);
  }

  // Flash LEDs on startup
  for(int j=0; j<5; j++) {
  for (int i = 0; i < ledCount; i++) digitalWrite(leds[i].pin, HIGH);
    delay(500);
    for (int i = 0; i < ledCount; i++) digitalWrite(leds[i].pin, LOW);
    delay(500);
  }

  // Connect to Wi-Fi
  WiFi.begin(ssid, password);
  Serial.print("Connecting WiFi");
  while (WiFi.status() != WL_CONNECTED) {
    Serial.print(".");
    delay(500);
  }
  Serial.println("\nWiFi connected!");

  // Setup MQTT
  mqtt.setServer(mqttHost, mqttPort);
  mqtt.setCallback(callback);
  connectMQTT();
  String msg=String(mqttClientID)+": Hello.";
  mqtt.publish(topic,msg.c_str());

}

// ===== Loop =====
void loop() {
  // Ensure MQTT is connected
  if (!mqtt.connected()) connectMQTT();
  mqtt.loop();

  // Periodic ping
  if (millis() - lastPing > pingInterval) {
    mqtt.publish(topic, "~ping~");
    lastPing = millis();
  }

  // Check Serial input (non-blocking)
  if (Serial.available() > 0) {
    String input = Serial.readStringUntil('\n');
    input.trim();
    if (input.length() > 0) {
      // Control LEDs locally
      for (int i = 0; i < ledCount; i++) {
        setLED(input, leds[i], false);
      }
      input=String(mqttClientID)+": "+input;
      // Publish message to MQTT
      mqtt.publish(topic, input.c_str());
      Serial.println("Published to chat: " + input);
    }
  }
}