#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"));
  }
}

// ===== 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);
  }

  // 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(300);
  }
  Serial.println("\nWiFi connected!");

  // Setup MQTT
  mqtt.setServer(mqttHost, mqttPort);
  mqtt.setCallback(callback);
  connectMQTT();
}

// ===== 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);
      }

      // Publish message to MQTT
      mqtt.publish(topic, input.c_str());
      Serial.println("Published to chat: " + input);
    }
  }
}