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

// WiFi credentials
const char* ssid = "Wokwi-GUEST";
const char* password = "";

// MQTT Broker settings
const char* mqtt_server = "public.cloud.shiftr.io";
const int mqtt_port = 1883;
const char* mqtt_user = "public";
const char* mqtt_password = "public";
const char* mqtt_topic = "oxigen-supergames2";

// LED pins
const int LED_GREEN = 2;   // LED hijau untuk "pake masker"
const int LED_RED = 4;     // LED merah untuk "tidak pake masker"
const int LED_BUILTIN = 2; // LED built-in untuk status koneksi

// WiFi and MQTT clients
WiFiClient espClient;
PubSubClient client(espClient);

// Variables
unsigned long lastMsg = 0;
float maskThreshold = 0.7;    // Threshold untuk deteksi masker
float noMaskThreshold = 0.7;  // Threshold untuk deteksi tidak pakai masker

void setup() {
  Serial.begin(115200);
  
  // Initialize LED pins
  pinMode(LED_GREEN, OUTPUT);
  pinMode(LED_RED, OUTPUT);
  pinMode(LED_BUILTIN, OUTPUT);
  
  // Turn off all LEDs initially
  digitalWrite(LED_GREEN, LOW);
  digitalWrite(LED_RED, LOW);
  digitalWrite(LED_BUILTIN, LOW);
  
  Serial.println("ESP32 Mask Detection System Starting...");
  
  // Setup WiFi
  setup_wifi();
  
  // Setup MQTT
  client.setServer(mqtt_server, mqtt_port);
  client.setCallback(callback);
  
  Serial.println("System ready!");
}

void setup_wifi() {
  delay(10);
  Serial.println();
  Serial.print("Connecting to WiFi: ");
  Serial.println(ssid);

  WiFi.begin(ssid, password);

  while (WiFi.status() != WL_CONNECTED) {
    delay(500);
    Serial.print(".");
    // Blink built-in LED while connecting
    digitalWrite(LED_BUILTIN, !digitalRead(LED_BUILTIN));
  }

  // WiFi connected - turn on built-in LED
  digitalWrite(LED_BUILTIN, HIGH);
  
  Serial.println();
  Serial.println("WiFi connected!");
  Serial.print("IP address: ");
  Serial.println(WiFi.localIP());
}

void callback(char* topic, byte* message, unsigned int length) {
  Serial.print("Message arrived on topic: ");
  Serial.println(topic);
  
  // Convert message to string
  String messageTemp;
  for (int i = 0; i < length; i++) {
    messageTemp += (char)message[i];
  }
  
  Serial.print("Message: ");
  Serial.println(messageTemp);
  
  // Parse the message untuk mendapatkan prediksi
  processPrediction(messageTemp);
}

void processPrediction(String message) {
  // Coba parse JSON: {"class":"pake masker","prob":0.85}
  String cls = "";
  float prob = NAN;

  message.trim();

  if (message.startsWith("{") && message.endsWith("}")) {
    // Parser JSON sangat sederhana (tanpa lib) untuk 2 field umum
    int iClass = message.indexOf("\"class\"");
    int iProb  = message.indexOf("\"prob\"");
    if (iClass >= 0) {
      int colon = message.indexOf(':', iClass);
      int q1 = message.indexOf('"', colon + 1);
      int q2 = message.indexOf('"', q1 + 1);
      if (q1 >= 0 && q2 > q1) {
        cls = message.substring(q1 + 1, q2);
      }
    }
    if (iProb >= 0) {
      int colon = message.indexOf(':', iProb);
      if (colon >= 0) {
        String s = message.substring(colon + 1);
        s.trim();
        // buang koma/kurung kurawal jika ada
        for (int i = 0; i < (int)s.length(); i++) {
          if (s[i] == ',' || s[i] == '}') { s = s.substring(0, i); break; }
        }
        prob = s.toFloat();
      }
    }
  }

  // Fallback "Class: prob"
  if (cls.length() == 0 || isnan(prob)) {
    int colonIndex = message.indexOf(':');
    if (colonIndex == -1) {
      Serial.println("Format pesan tidak dikenali");
      return;
    }
    cls = message.substring(0, colonIndex);
    String probabilityStr = message.substring(colonIndex + 1);
    probabilityStr.trim();
    prob = probabilityStr.toFloat();
  }

  // Normalisasi teks kelas
  cls.trim();
  cls.toLowerCase();

  // --- KLASIFIKASI YANG TEPAT ---
  // Deteksi "no mask" / "tidak" lebih dulu agar tidak tertangkap oleh cek "mask"
  bool isNoMask =
    (cls.indexOf("no mask") >= 0) ||
    (cls.indexOf("no_mask") >= 0) ||
    (cls.indexOf("nomask") >= 0) ||
    (cls.indexOf("tidak") >= 0) ||
    (cls.indexOf("tanpa") >= 0);

  bool isMask = (!isNoMask) && (
    (cls.indexOf("pake masker") >= 0) ||
    (cls.indexOf("pakai masker") >= 0) ||
    (cls == "mask") ||
    (cls.indexOf(" with mask") >= 0) ||       // jika format English
    (cls.indexOf("mask") >= 0)                 // terakhir, aman karena isNoMask sudah ditentukan duluan
  );

  Serial.println("Processed prediction:");
  Serial.println("Class: " + cls);
  Serial.println("Probability: " + String(prob, 3));
  Serial.print("isMask=");   Serial.print(isMask);
  Serial.print(" isNoMask=");Serial.println(isNoMask);

  // Reset dulu
  digitalWrite(LED_GREEN, LOW);
  digitalWrite(LED_RED, LOW);

  // Terapkan ambang per kelas
  if (isMask && prob >= maskThreshold) {
    digitalWrite(LED_GREEN, HIGH);
    Serial.println("✅ MASK DETECTED - Green LED ON");
    blinkLED(LED_GREEN, 2, 150);
  } else if (isNoMask && prob >= noMaskThreshold) {
    digitalWrite(LED_RED, HIGH);
    Serial.println("❌ NO MASK DETECTED - Red LED ON");
    blinkLED(LED_RED, 3, 120);
  } else {
    // Confidence tidak cukup → semua mati
    Serial.println("⚠️ Confidence rendah / kelas tidak dikenal - LEDs OFF");
  }
}

void blinkLED(int ledPin, int times, int delayMs) {
  for (int i = 0; i < times; i++) {
    digitalWrite(ledPin, LOW);
    delay(delayMs);
    digitalWrite(ledPin, HIGH);
    delay(delayMs);
  }
}

void reconnect() {
  // Loop until we're reconnected
  while (!client.connected()) {
    Serial.print("Attempting MQTT connection...");
    
    // Create a random client ID
    String clientId = "ESP32Client-";
    clientId += String(random(0xffff), HEX);
    
    // Attempt to connect
    if (client.connect(clientId.c_str(), mqtt_user, mqtt_password)) {
      Serial.println("connected");
      
      // Subscribe to the topic
      client.subscribe(mqtt_topic);
      Serial.print("Subscribed to topic: ");
      Serial.println(mqtt_topic);
      
      // Turn on built-in LED to indicate MQTT connection
      digitalWrite(LED_BUILTIN, HIGH);
      
    } else {
      Serial.print("failed, rc=");
      Serial.print(client.state());
      Serial.println(" try again in 5 seconds");
      
      // Blink built-in LED to indicate connection problem
      for (int i = 0; i < 10; i++) {
        digitalWrite(LED_BUILTIN, !digitalRead(LED_BUILTIN));
        delay(500);
      }
    }
  }
}

void loop() {
  // Check WiFi connection
  if (WiFi.status() != WL_CONNECTED) {
    Serial.println("WiFi connection lost! Reconnecting...");
    setup_wifi();
  }
  
  // Check MQTT connection
  if (!client.connected()) {
    digitalWrite(LED_BUILTIN, LOW);
    reconnect();
  }
  
  client.loop();
  
  // Send heartbeat every 30 seconds
  unsigned long now = millis();
  if (now - lastMsg > 30000) {
    lastMsg = now;
    Serial.println("System running... Waiting for predictions...");
    
    // Quick blink to show system is alive
    digitalWrite(LED_BUILTIN, LOW);
    delay(50);
    digitalWrite(LED_BUILTIN, HIGH);
  }
  
  // Small delay to prevent excessive CPU usage
  delay(100);
}

// Fungsi tambahan untuk testing LED
void testLEDs() {
  Serial.println("Testing LEDs...");
  
  // Test Green LED
  Serial.println("Green LED ON");
  digitalWrite(LED_GREEN, HIGH);
  delay(1000);
  digitalWrite(LED_GREEN, LOW);
  
  // Test Red LED  
  Serial.println("Red LED ON");
  digitalWrite(LED_RED, HIGH);
  delay(1000);
  digitalWrite(LED_RED, LOW);
  
  // Test both LEDs
  Serial.println("Both LEDs ON");
  digitalWrite(LED_GREEN, HIGH);
  digitalWrite(LED_RED, HIGH);
  delay(1000);
  digitalWrite(LED_GREEN, LOW);
  digitalWrite(LED_RED, LOW);
  
  Serial.println("LED test complete");
}