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

// ═══ NODE IDENTITY — change only this for Node B / C ═══
const char* NODE_ID  = "nodeA";   // → "nodeB" or "nodeC"
const int   NODE_NUM = 1;          // → 2 or 3

// ═══ WIFI — Wokwi provides this automatically ═══
const char* WIFI_SSID = "Wokwi-GUEST";
const char* WIFI_PASS = "";

// ═══ MQTT ═══
const char* MQTT_BROKER = "broker.hivemq.com";
const int   MQTT_PORT   = 1883;

// ═══ GPIO PINS — matched exactly to your diagram.json ═══
// dht1:SDA  → esp:15
// led2 (GREEN)  → esp:2  via r1 (220Ω)  ← HEALTHY
// led3 (YELLOW) → esp:4  via r2 (220Ω)  ← WARNING
// led1 (RED)    → esp:5  via r3 (220Ω)  ← CRITICAL
// All LED cathodes → esp:CMD (GND)
#define DHT_PIN     15
#define LED_GREEN    2   // led2 — green
#define LED_YELLOW   4   // led3 — yellow
#define LED_RED      5   // led1 — red

// ═══ SIMULATED NETWORK METRICS ═══
float battery  = 100.0;
float rssi     = -55.0;
float latency  =  40.0;
float pdr      =  98.0;
float healthScore = 100.0;

// Moving average buffer for anomaly detection
#define MA_SIZE 5
float  tempBuf[MA_SIZE] = {0};
int    bufIdx  = 0;
bool   bufFull = false;

// Adaptive TX interval
unsigned long txInterval = 5000;
unsigned long lastTx     = 0;

DHTesp        dht;
WiFiClient    wifiClient;
PubSubClient  mqtt(wifiClient);

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

  pinMode(LED_GREEN,  OUTPUT);
  pinMode(LED_YELLOW, OUTPUT);
  pinMode(LED_RED,    OUTPUT);

  // Startup blink — all 3 LEDs flash once to confirm wiring
  digitalWrite(LED_GREEN,  HIGH);
  digitalWrite(LED_YELLOW, HIGH);
  digitalWrite(LED_RED,    HIGH);
  delay(600);
  digitalWrite(LED_GREEN,  LOW);
  digitalWrite(LED_YELLOW, LOW);
  digitalWrite(LED_RED,    LOW);

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

  // WiFi
  Serial.printf("\n[%s] Connecting WiFi", NODE_ID);
  WiFi.begin(WIFI_SSID, WIFI_PASS);
  while (WiFi.status() != WL_CONNECTED) {
    delay(400);
    Serial.print(".");
  }
  Serial.printf(" OK  IP: %s\n", WiFi.localIP().toString().c_str());

  // MQTT
  mqtt.setServer(MQTT_BROKER, MQTT_PORT);
  mqtt.setCallback(onMqttMessage);
  connectMQTT();

  Serial.printf("[%s] Node online. Publishing every %lums\n", NODE_ID, txInterval);
}

// ════════════════════════════════════════
void loop() {
  if (!mqtt.connected()) connectMQTT();
  mqtt.loop();

  unsigned long now = millis();
  if (now - lastTx >= txInterval) {
    lastTx = now;
    simulateDegradation();
    readAndPublish();
  }
}

// ════════════════════════════════════════
// Simulate realistic IoT node aging over time
void simulateDegradation() {
  // Battery drains slowly (0.08% per TX cycle)
  battery -= 0.08f;
  battery  = max(0.0f, battery);

  // RSSI fluctuates ±3 dBm
  rssi += (float)random(-3, 4);
  rssi  = constrain(rssi, -95.0f, -30.0f);

  // Latency creeps when battery is low
  latency += (battery < 30.0f) ? (float)random(1, 8)
                                : (float)random(-2, 3);
  latency  = constrain(latency, 10.0f, 300.0f);

  // PDR degrades with low battery or high latency
  if (battery < 20.0f || latency > 150.0f)
    pdr -= (float)random(0, 3);
  else
    pdr += (float)random(0, 2);
  pdr = constrain(pdr, 0.0f, 100.0f);
}

// ════════════════════════════════════════
// Weighted health score (max = 100)
float computeHealth() {
  float latScore = max(0.0f, 100.0f - latency / 2.0f);
  return (battery * 0.40f) + (pdr * 0.35f) + (latScore * 0.25f);
}

String getState(float score) {
  if (score >= 70.0f) return "HEALTHY";
  if (score >= 40.0f) return "WARNING";
  return "CRITICAL";
}

// ════════════════════════════════════════
// Moving-average anomaly detection on temperature
bool detectAnomaly(float temp) {
  tempBuf[bufIdx] = temp;
  bufIdx = (bufIdx + 1) % MA_SIZE;
  if (bufIdx == 0) bufFull = true;

  int count = bufFull ? MA_SIZE : bufIdx;
  if (count < 2) return false;

  float sum = 0;
  for (int i = 0; i < count; i++) sum += tempBuf[i];
  float avg = sum / count;

  return (abs(temp - avg) > 5.0f);   // anomaly if >5°C from rolling avg
}

// ════════════════════════════════════════
// Update the 3 status LEDs
void updateLEDs(const String& state) {
  digitalWrite(LED_GREEN,  state == "HEALTHY"  ? HIGH : LOW);
  digitalWrite(LED_YELLOW, state == "WARNING"  ? HIGH : LOW);
  digitalWrite(LED_RED,    state == "CRITICAL" ? HIGH : LOW);
}

// Adaptive transmission interval
void adaptInterval(const String& state) {
  if (state == "HEALTHY")  txInterval = 8000;   // low-power mode
  if (state == "WARNING")  txInterval = 4000;   // normal mode
  if (state == "CRITICAL") txInterval = 1500;   // high-frequency mode
}

// ════════════════════════════════════════
// Main read + publish cycle
void readAndPublish() {
  TempAndHumidity data = dht.getTempAndHumidity();
  float temp = isnan(data.temperature) ? 25.0f : data.temperature;
  float hum  = isnan(data.humidity)    ? 50.0f : data.humidity;

  healthScore    = computeHealth();
  String state   = getState(healthScore);
  bool   anomaly = detectAnomaly(temp);

  updateLEDs(state);
  adaptInterval(state);

  // ── Sensor data payload ──
  StaticJsonDocument<256> dDoc;
  dDoc["node"]     = NODE_ID;
  dDoc["temp"]     = round(temp  * 10) / 10.0;
  dDoc["humidity"] = round(hum   * 10) / 10.0;
  dDoc["battery"]  = round(battery * 10) / 10.0;
  dDoc["rssi"]     = (int)rssi;
  dDoc["latency"]  = (int)latency;
  dDoc["pdr"]      = round(pdr * 10) / 10.0;
  dDoc["anomaly"]  = anomaly ? 1 : 0;
  dDoc["ts"]       = millis();

  char dBuf[256];
  serializeJson(dDoc, dBuf);

  // ── Health payload ──
  StaticJsonDocument<128> hDoc;
  hDoc["node"]  = NODE_ID;
  hDoc["score"] = round(healthScore * 10) / 10.0;
  hDoc["state"] = state;
  hDoc["txInt"] = (int)txInterval;

  char hBuf[128];
  serializeJson(hDoc, hBuf);

  // ── MQTT publish ──
  char dataTopic[32], healthTopic[32];
  sprintf(dataTopic,   "iot/%s/data",   NODE_ID);
  sprintf(healthTopic, "iot/%s/health", NODE_ID);

  mqtt.publish(dataTopic,   dBuf);
  mqtt.publish(healthTopic, hBuf);

  // ── Serial Monitor output ──
  Serial.printf("[%s] %s | Score:%.1f Bat:%.1f%% Temp:%.1f°C Hum:%.1f%% "
                "PDR:%.1f%% Lat:%.0fms Anomaly:%s\n",
    NODE_ID, state.c_str(), healthScore,
    battery, temp, hum, pdr, latency,
    anomaly ? "YES!" : "no");
}

// ════════════════════════════════════════
// Receive self-healing commands from gateway
void onMqttMessage(char* topic, byte* payload, unsigned int len) {
  char msg[128];
  memcpy(msg, payload, min(len, (unsigned int)127));
  msg[min(len, (unsigned int)127)] = '\0';

  StaticJsonDocument<128> doc;
  if (deserializeJson(doc, msg) != DeserializationError::Ok) return;

  String action = doc["action"] | "";

  if (action == "reduce_freq") {
    txInterval = 15000;
    Serial.printf("[%s][CMD] reduce_freq → TX every 15s\n", NODE_ID);
  } else if (action == "increase_freq") {
    txInterval = 1000;
    Serial.printf("[%s][CMD] increase_freq → TX every 1s\n", NODE_ID);
  } else if (action == "ping") {
    Serial.printf("[%s][CMD] ping from gateway\n", NODE_ID);
  }
}

// ════════════════════════════════════════
void connectMQTT() {
  while (!mqtt.connected()) {
    char clientId[40];
    sprintf(clientId, "wokwi-%s-%d", NODE_ID, random(9999));
    Serial.printf("[%s] MQTT connecting as %s ...\n", NODE_ID, clientId);

    if (mqtt.connect(clientId)) {
      char cmdTopic[32];
      sprintf(cmdTopic, "iot/%s/cmd", NODE_ID);
      mqtt.subscribe(cmdTopic);
      Serial.printf("[%s] MQTT connected. Subscribed to %s\n", NODE_ID, cmdTopic);
    } else {
      Serial.printf("[%s] MQTT failed rc=%d — retry in 3s\n",
                    NODE_ID, mqtt.state());
      delay(3000);
    }
  }
}