/*
  CYD Lightning Detector
  Supports both ILI9341 and ST7789 CYD display variants.

  Required libraries:
    - GFX Library for Arduino
    - DFRobot_AS3935
    - PubSubClient by Nick O'Leary

  Home Assistant / MQTT behavior:
    - Display/backlight remains ON continuously.
    - Wi-Fi is OFF during normal AS3935 monitoring.
    - At startup, Wi-Fi is briefly enabled with AS3935 IRQ disabled to
      publish retained Home Assistant MQTT Discovery.
    - Only the FIRST legitimate lightning strike in each 30-minute strike
      session turns Wi-Fi on and publishes one MQTT "lightning" event.
    - Each alert includes an ack_id. Home Assistant must echo that ack_id to
      lightning/ack. The detector waits for confirmation and retries up to
      three times during the same controlled Wi-Fi window.
    - Subsequent strikes update the local display/count only.
    - The 30-minute session timer is restarted by each legitimate strike.
    - After 30 minutes with no lightning, the next strike may alert again.
    - AS3935 IRQ is disabled while Wi-Fi is active, then Wi-Fi is shut
      completely off, a 2-second settling delay is observed, pending AS3935
      interrupt state is cleared, and monitoring resumes.

  Select ONE display type below.
*/

// ============================================================
// Display selection
// ============================================================

// Original CYD boards:
#define DISPLAY_ILI9341

// Newer CYD boards:
//#define DISPLAY_ST7789

#if defined(DISPLAY_ILI9341) && defined(DISPLAY_ST7789)
  #error "Select only one display type."
#endif

#if !defined(DISPLAY_ILI9341) && !defined(DISPLAY_ST7789)
  #error "Select a display type."
#endif

#include <Wire.h>
#include <SPI.h>
#include <Arduino_GFX_Library.h>
#include <DFRobot_AS3935_I2C.h>
#include <WiFi.h>
#include <PubSubClient.h>

// ============================================================
// CYD display pins
// ============================================================
#define TFT_MOSI 13
#define TFT_MISO 12
#define TFT_CLK  14
#define TFT_CS   15
#define TFT_DC   2
#define TFT_RST  GFX_NOT_DEFINED
#define TFT_BL   21

// ============================================================
// Display colors (RGB565)
// ============================================================
#define COLOR_BLACK     0x0000
#define COLOR_BLUE      0x001F
#define COLOR_RED       0xF800
#define COLOR_GREEN     0x07E0
#define COLOR_CYAN      0x07FF
#define COLOR_MAGENTA   0xF81F
#define COLOR_YELLOW    0xFFE0
#define COLOR_WHITE     0xFFFF
#define COLOR_ORANGE    0xFD20
#define COLOR_NAVY      0x000F
#define COLOR_DARKCYAN  0x03EF

// ============================================================
// Display bus and driver
// ============================================================
Arduino_DataBus *bus = new Arduino_ESP32SPI(
  TFT_DC,
  TFT_CS,
  TFT_CLK,
  TFT_MOSI,
  TFT_MISO,
  HSPI
);

#ifdef DISPLAY_ILI9341
Arduino_GFX *gfx = new Arduino_ILI9341(
  bus,
  TFT_RST,
  3
);
const char *displayName = "ILI9341";
#endif

#ifdef DISPLAY_ST7789
Arduino_GFX *gfx = new Arduino_ST7789(
  bus,
  TFT_RST,
  3,
  false
);
const char *displayName = "ST7789";
#endif

// ============================================================
// AS3935 pins
// ============================================================
#define AS3935_SDA 27
#define AS3935_SCL 22
#define AS3935_IRQ 35
#define AS3935_ADDR AS3935_ADD3

DFRobot_AS3935_I2C lightning(AS3935_IRQ, AS3935_ADDR);

// ============================================================
// Buzzer
// ============================================================
#define BUZZER_PIN 26
const int BUZZER_RESOLUTION = 8;

// ============================================================
// Wi-Fi / MQTT configuration
// ============================================================
// Enter your Wi-Fi credentials and the password for the dedicated
// Home Assistant local user "lightning".
const char *WIFI_SSID     = "101";
const char *WIFI_PASSWORD = "Samsung123456";

const char *MQTT_HOST     = "192.168.2.118";
const uint16_t MQTT_PORT  = 1883;
const char *MQTT_USER     = "lightning";
const char *MQTT_PASSWORD = "lightning";

const char *MQTT_CLIENT_ID = "cyd_lightning_detector";
const char *TOPIC_EVENT    = "lightning/event";
const char *TOPIC_ACK      = "lightning/ack";

#define MQTT_ACK_WAIT_MS        2500UL
#define MQTT_ALERT_MAX_ATTEMPTS 3
#define MQTT_RETRY_DELAY_MS     500UL

// Home Assistant discovery topic for one MQTT Event entity.
const char *HA_DISCOVERY_TOPIC =
  "homeassistant/event/cyd_lightning/detector_event/config";

WiFiClient wifiClient;
PubSubClient mqtt(wifiClient);

volatile bool mqttAckReceived = false;
String expectedAckId;

// ============================================================
// Timeouts
// ============================================================
#define STRIKE_RESET_TIMEOUT_MS 1800000UL  // 30 minutes after LAST legitimate strike
#define EVENT_STATUS_TIMEOUT_MS 5000UL     // 5 seconds

// Wi-Fi/MQTT is normally OFF to prevent ESP32 RF from disturbing the AS3935.
// It is enabled only at startup (while AS3935 IRQ is disabled) and for the
// first legitimate lightning strike in each 30-minute strike session.
#define WIFI_CONNECT_TIMEOUT_MS 10000UL
#define MQTT_CONNECT_TIMEOUT_MS 5000UL
#define WIFI_SETTLE_AFTER_OFF_MS 2000UL

// ============================================================
// Runtime state
// ============================================================
volatile bool irqFlag = false;

uint32_t strikeCount = 0;
uint8_t lastDistanceKm = 0;
uint32_t lastEnergy = 0;

String lastEvent = "Startup";

unsigned long lastEventMillis = 0;
unsigned long lastRefreshMillis = 0;
unsigned long lastStrikeMillis = 0;

// True after the first MQTT alert in the current strike session.
// It is reset only after 30 minutes with no legitimate lightning strike.
bool mqttAlertSentThisSession = false;

// ============================================================
// IRQ handler
// ============================================================
void IRAM_ATTR onLightningIRQ() {
  irqFlag = true;
}

// ============================================================
// Buzzer helpers
// ============================================================
void initBuzzer() {
  ledcAttach(BUZZER_PIN, 3500, BUZZER_RESOLUTION);
  ledcWriteTone(BUZZER_PIN, 0);
}

void beepLightning() {
  ledcWriteTone(BUZZER_PIN, 3500);
  delay(400);
  ledcWriteTone(BUZZER_PIN, 0);
}

// ============================================================
// Display helpers
// ============================================================
void clearValueRow(int y) {
  gfx->fillRect(135, y, gfx->width() - 150, 22, COLOR_BLACK);
}

String ageString() {
  if (lastEventMillis == 0) {
    return "--";
  }

  unsigned long ageSec = (millis() - lastEventMillis) / 1000UL;

  char buf[20];
  snprintf(buf, sizeof(buf), "%lus", ageSec);

  return String(buf);
}

void drawHeader() {
  gfx->fillScreen(COLOR_BLACK);

  gfx->setTextColor(COLOR_YELLOW, COLOR_BLACK);
  gfx->setTextSize(2);
  gfx->setCursor(12, 8);
  gfx->print("AS3935 Lightning");

  gfx->drawFastHLine(
    0,
    32,
    gfx->width(),
    COLOR_DARKCYAN
  );
}

void drawLabels() {
  gfx->setTextSize(2);
  gfx->setTextColor(COLOR_WHITE, COLOR_BLACK);

  gfx->setCursor(10, 45);
  gfx->print("Status:");

  gfx->setCursor(10, 75);
  gfx->print("Strikes:");

  gfx->setCursor(10, 105);
  gfx->print("Distance:");

  gfx->setCursor(10, 135);
  gfx->print("Energy:");

  gfx->setCursor(10, 165);
  gfx->print("Noise:");

  gfx->setCursor(10, 195);
  gfx->print("Last age:");
}

void drawFooter() {
  const int footerHeight = 15;
  const int footerY = gfx->height() - footerHeight;

  gfx->fillRect(
    0,
    footerY,
    gfx->width(),
    footerHeight,
    COLOR_NAVY
  );

  gfx->setTextSize(1);
  gfx->setTextColor(COLOR_WHITE, COLOR_NAVY);
  gfx->setCursor(6, footerY + 4);
  gfx->print("IRQ=35 SDA=27 SCL=22");
}

void drawValues() {
  clearValueRow(45);
  clearValueRow(75);
  clearValueRow(105);
  clearValueRow(135);
  clearValueRow(165);
  clearValueRow(195);

  gfx->setTextSize(2);

  uint16_t statusColor = COLOR_GREEN;

  if (lastEvent == "Lightning") {
    statusColor = COLOR_RED;
  } else if (lastEvent == "Disturber") {
    statusColor = COLOR_ORANGE;
  } else if (lastEvent == "Noise high") {
    statusColor = COLOR_MAGENTA;
  } else if (lastEvent == "Init fail") {
    statusColor = COLOR_RED;
  }

  gfx->setTextColor(statusColor, COLOR_BLACK);
  gfx->setCursor(135, 45);
  gfx->print(lastEvent);

  gfx->setTextColor(COLOR_CYAN, COLOR_BLACK);
  gfx->setCursor(135, 75);
  gfx->print(strikeCount);

  gfx->setCursor(135, 105);

  if (lastDistanceKm == 0) {
    gfx->setTextColor(COLOR_WHITE, COLOR_BLACK);
    gfx->print("--");
  } else {
    float miles = lastDistanceKm * 0.621371f;

    uint16_t distColor;

    if (miles <= 10.0f) {
      distColor = COLOR_RED;
    } else if (miles <= 20.0f) {
      distColor = COLOR_YELLOW;
    } else {
      distColor = COLOR_GREEN;
    }

    gfx->setTextColor(distColor, COLOR_BLACK);
    gfx->print(miles, 1);
    gfx->print(" mi");
  }

  gfx->setTextColor(COLOR_WHITE, COLOR_BLACK);
  gfx->setCursor(135, 135);

  if (lastEnergy == 0) {
    gfx->print("--");
  } else {
    gfx->print(lastEnergy);
  }

  gfx->setCursor(135, 165);
  gfx->print(lightning.getNoiseFloorLvl());

  gfx->setCursor(135, 195);
  gfx->print(ageString());
}

void drawAll() {
  drawHeader();
  drawLabels();
  drawValues();
  drawFooter();
}

// ============================================================
// MQTT callback / ACK handling
// ============================================================

void mqttCallback(char *topic, byte *payload, unsigned int length) {
  if (String(topic) != TOPIC_ACK) {
    return;
  }

  String message;
  message.reserve(length);

  for (unsigned int i = 0; i < length; i++) {
    message += (char)payload[i];
  }

  message.trim();

  if (message == expectedAckId) {
    mqttAckReceived = true;
    Serial.print("MQTT ACK received: ");
    Serial.println(message);
  }
}

// ============================================================
// Wi-Fi / MQTT burst helpers
// ============================================================

void detachLightningIRQ() {
  detachInterrupt(digitalPinToInterrupt(AS3935_IRQ));

  noInterrupts();
  irqFlag = false;
  interrupts();
}

void attachLightningIRQ() {
  noInterrupts();
  irqFlag = false;
  interrupts();

  attachInterrupt(
    digitalPinToInterrupt(AS3935_IRQ),
    onLightningIRQ,
    RISING
  );
}

void turnWiFiOffAndSettle() {
  if (mqtt.connected()) {
    mqtt.disconnect();
  }

  wifiClient.stop();

  // Turn the ESP32 radio completely off. Do not erase saved credentials.
  WiFi.disconnect(true, false);
  WiFi.mode(WIFI_OFF);

  // Allow RF/power-supply transients to disappear before listening again.
  delay(WIFI_SETTLE_AFTER_OFF_MS);

  // Clear any AS3935 interrupt condition caused while Wi-Fi was active.
  // Reading the interrupt source clears the pending source indication.
  (void)lightning.getInterruptSrc();
  delay(10);
  (void)lightning.getInterruptSrc();

  noInterrupts();
  irqFlag = false;
  interrupts();
}

bool connectWiFiBurst() {
  WiFi.mode(WIFI_STA);

  // Lowest practical ESP32 Wi-Fi TX power. The radio is still kept OFF
  // during normal monitoring; this merely reduces the burst power further.
  WiFi.setTxPower(WIFI_POWER_2dBm);

  WiFi.setAutoReconnect(false);
  WiFi.persistent(false);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);

  Serial.print("Wi-Fi burst: connecting");

  unsigned long started = millis();

  while (WiFi.status() != WL_CONNECTED) {
    if (millis() - started >= WIFI_CONNECT_TIMEOUT_MS) {
      Serial.println(" timeout");
      return false;
    }

    delay(100);
    Serial.print(".");
  }

  Serial.println();
  Serial.print("Wi-Fi connected. IP: ");
  Serial.println(WiFi.localIP());

  return true;
}

bool connectMQTTBurst() {
  mqtt.setServer(MQTT_HOST, MQTT_PORT);
  mqtt.setCallback(mqttCallback);
  mqtt.setBufferSize(1024);
  mqtt.setKeepAlive(15);
  mqtt.setSocketTimeout(3);

  unsigned long started = millis();

  while (!mqtt.connected()) {
    if (mqtt.connect(MQTT_CLIENT_ID, MQTT_USER, MQTT_PASSWORD)) {
      Serial.println("MQTT connected.");
      mqtt.subscribe(TOPIC_ACK);
      mqtt.loop();
      return true;
    }

    if (millis() - started >= MQTT_CONNECT_TIMEOUT_MS) {
      Serial.print("MQTT timeout; state=");
      Serial.println(mqtt.state());
      return false;
    }

    delay(250);
  }

  return true;
}

void publishHomeAssistantDiscovery() {
  // Retained discovery config remains in Mosquitto while the detector's
  // Wi-Fi is off, so Home Assistant keeps the entity available for events.
  String payload = "{";
  payload += "\"name\":\"Lightning alert\",";
  payload += "\"unique_id\":\"cyd_lightning_first_strike_event\",";
  payload += "\"state_topic\":\"";
  payload += TOPIC_EVENT;
  payload += "\",";
  payload += "\"event_types\":[\"lightning\"],";
  payload += "\"device\":{";
  payload += "\"identifiers\":[\"cyd_as3935_lightning\"],";
  payload += "\"name\":\"CYD Lightning Detector\",";
  payload += "\"manufacturer\":\"DIY\",";
  payload += "\"model\":\"ESP32 CYD + DFRobot AS3935\",";
  payload += "\"sw_version\":\"First-strike MQTT\"";
  payload += "}";
  payload += "}";

  mqtt.publish(HA_DISCOVERY_TOPIC, payload.c_str(), true);

  // Give the broker/HA a moment to process discovery before any first event.
  for (int i = 0; i < 5; i++) {
    mqtt.loop();
    delay(100);
  }

  Serial.println("Home Assistant MQTT discovery published.");
}

bool sendFirstStrikeMQTTAlert(uint8_t distanceKm, uint32_t energy) {
  // IMPORTANT:
  // The AS3935 interrupt is already detached before entering this function.
  // Wi-Fi interference therefore cannot create another accepted strike while
  // the radio is operating.

  bool delivered = false;

  if (connectWiFiBurst()) {
    if (connectMQTTBurst()) {
      // Republish retained discovery each burst. This ensures the HA entity
      // remains available after broker/HA maintenance.
      publishHomeAssistantDiscovery();

      float distanceMi = distanceKm * 0.621371f;

      // Unique ACK token for this one alert. Home Assistant must echo this
      // exact value back to lightning/ack after receiving the event.
      expectedAckId =
        String("strike-") +
        String(strikeCount) + "-" +
        String((unsigned long)millis());

      String payload = "{";
      payload += "\"event_type\":\"lightning\",";
      payload += "\"ack_id\":\"";
      payload += expectedAckId;
      payload += "\",";
      payload += "\"distance_km\":";
      payload += String(distanceKm);
      payload += ",";
      payload += "\"distance_mi\":";
      payload += String(distanceMi, 1);
      payload += ",";
      payload += "\"energy\":";
      payload += String(energy);
      payload += ",";
      payload += "\"strike_count\":";
      payload += String(strikeCount);
      payload += "}";

      for (int attempt = 1;
           attempt <= MQTT_ALERT_MAX_ATTEMPTS && !delivered;
           attempt++) {

        mqttAckReceived = false;

        Serial.print("MQTT alert attempt ");
        Serial.print(attempt);
        Serial.print("/");
        Serial.println(MQTT_ALERT_MAX_ATTEMPTS);

        bool publishAccepted =
          mqtt.publish(TOPIC_EVENT, payload.c_str(), false);

        if (!publishAccepted) {
          Serial.println("MQTT publish call failed.");
        }
        else {
          unsigned long ackStart = millis();

          while (!mqttAckReceived &&
                 millis() - ackStart < MQTT_ACK_WAIT_MS) {

            if (!mqtt.connected()) {
              Serial.println("MQTT disconnected while waiting for ACK.");
              break;
            }

            mqtt.loop();
            delay(10);
          }

          if (mqttAckReceived) {
            delivered = true;
            Serial.println("Lightning alert delivery confirmed by Home Assistant.");
            break;
          }

          Serial.println("No Home Assistant ACK received.");
        }

        if (attempt < MQTT_ALERT_MAX_ATTEMPTS) {
          delay(MQTT_RETRY_DELAY_MS);

          // Reconnect MQTT if the broker dropped during an attempt.
          if (!mqtt.connected()) {
            Serial.println("Reconnecting MQTT for retry...");
            if (!connectMQTTBurst()) {
              Serial.println("MQTT reconnect failed.");
              break;
            }
          }
        }
      }
    }
  }

  if (!delivered) {
    Serial.println("WARNING: lightning alert was not ACK-confirmed.");
  }

  // Whether delivery succeeded or failed, restore the detector to the
  // RF-quiet monitoring state.
  turnWiFiOffAndSettle();

  return delivered;
}

void registerMQTTDiscoveryAtStartup() {
  // At startup the AS3935 IRQ is deliberately NOT attached yet, so any
  // interference created by Wi-Fi cannot be interpreted as lightning.
  Serial.println("Registering Home Assistant MQTT entity...");

  if (connectWiFiBurst() && connectMQTTBurst()) {
    publishHomeAssistantDiscovery();
  }

  turnWiFiOffAndSettle();
}

// ============================================================
// State reset helpers
// ============================================================
void resetStrikeSession() {
  strikeCount = 0;
  lastDistanceKm = 0;
  lastEnergy = 0;

  lastEvent = "Monitoring";
  lastEventMillis = millis();
  lastStrikeMillis = 0;
  mqttAlertSentThisSession = false;

  drawValues();
}

void resetTransientStatusIfTimedOut() {
  bool transientStatus =
    lastEvent == "Lightning" ||
    lastEvent == "Disturber" ||
    lastEvent == "Noise high";

  if (transientStatus &&
      millis() - lastEventMillis >= EVENT_STATUS_TIMEOUT_MS) {

    lastEvent = "Monitoring";
    lastEventMillis = millis();

    drawValues();
  }
}

void resetStrikeDisplayIfTimedOut() {
  if (strikeCount == 0 || lastStrikeMillis == 0) {
    return;
  }

  if (millis() - lastStrikeMillis >= STRIKE_RESET_TIMEOUT_MS) {
    resetStrikeSession();
  }
}

// ============================================================
// Sensor initialization
// ============================================================
bool initSensor() {
  Wire.begin(AS3935_SDA, AS3935_SCL);
  delay(50);

  if (lightning.begin() != 0) {
    return false;
  }

  lightning.setOutdoors();
  lightning.disturberDis();
  lightning.setNoiseFloorLvl(1);
  lightning.setWatchdogThreshold(1);
  lightning.setSpikeRejection(1);

  return true;
}

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

  pinMode(TFT_BL, OUTPUT);
  // Line-powered application: keep display/backlight on continuously.
  digitalWrite(TFT_BL, HIGH);

  if (!gfx->begin()) {
    Serial.println("Display initialization failed.");

    while (true) {
      delay(1000);
    }
  }

  Serial.println();
  Serial.println("CYD Lightning Detector");
  Serial.print("Display driver: ");
  Serial.println(displayName);
  Serial.print("Display size: ");
  Serial.print(gfx->width());
  Serial.print(" x ");
  Serial.println(gfx->height());

  initBuzzer();

  gfx->fillScreen(COLOR_BLACK);
  gfx->setTextColor(COLOR_WHITE, COLOR_BLACK);
  gfx->setTextSize(2);
  gfx->setCursor(10, 60);
  gfx->print("Initializing...");

  pinMode(AS3935_IRQ, INPUT);

  // Do NOT attach the AS3935 interrupt yet. Wi-Fi is used briefly during
  // startup to register MQTT Discovery, and Wi-Fi is known to interfere
  // with this detector hardware.
  if (!initSensor()) {
    lastEvent = "Init fail";
    lastEventMillis = millis();

    drawAll();

    while (true) {
      delay(1000);
    }
  }

  lastEvent = "Monitoring";
  lastEventMillis = millis();

  drawAll();

  // Brief startup network burst while AS3935 IRQ is disabled.
  // This creates the Home Assistant MQTT Event entity without allowing
  // startup Wi-Fi RF to register as a lightning strike.
  registerMQTTDiscoveryAtStartup();

  // Normal operating state begins here: Wi-Fi OFF, display ON, AS3935 armed.
  attachLightningIRQ();

  Serial.println("Monitoring: Wi-Fi OFF, AS3935 armed.");
}

// ============================================================
// Event handling
// ============================================================
void handleLightningInterrupt() {
  uint8_t src = lightning.getInterruptSrc();

  if (src == 1) {
    strikeCount++;

    lastDistanceKm = lightning.getLightningDistKm();

    uint32_t energy = lightning.getStrikeEnergyRaw();

    if (energy > 0) {
      lastEnergy = energy;
    }

    lastEvent = "Lightning";
    lastEventMillis = millis();
    lastStrikeMillis = millis();

    beepLightning();

    Serial.println();
    Serial.println("LIGHTNING DETECTED");
    Serial.print("Strike count: ");
    Serial.println(strikeCount);
    Serial.print("Distance (km): ");
    Serial.println(lastDistanceKm);
    Serial.print("Energy raw: ");
    Serial.println(energy);

    // Only the first legitimate strike in the current session causes
    // a network transmission. Every legitimate strike still updates
    // lastStrikeMillis, so the session ends only after 30 quiet minutes.
    if (!mqttAlertSentThisSession) {
      // Latch BEFORE attempting the network connection so a failed Wi-Fi
      // attempt cannot result in repeated RF bursts during this same storm.
      mqttAlertSentThisSession = true;

      // Prevent ESP32 Wi-Fi activity from generating accepted AS3935 IRQs.
      detachLightningIRQ();

      Serial.println("First strike in session: sending one MQTT alert.");
      (void)sendFirstStrikeMQTTAlert(lastDistanceKm, energy);

      // Wi-Fi is now OFF and the AS3935 has had time to settle/clear.
      attachLightningIRQ();

      Serial.println("AS3935 monitoring resumed; Wi-Fi OFF.");
    }
  }
  else if (src == 2) {
    lastEvent = "Disturber";
    lastEventMillis = millis();

    Serial.println("Disturber detected.");
  }
  else if (src == 3) {
    lastEvent = "Noise high";
    lastEventMillis = millis();

    Serial.println("Noise level too high.");
  }
  else {
    Serial.print("Unknown interrupt source: ");
    Serial.println(src);
  }

  drawValues();
}

// ============================================================
// Main loop
// ============================================================
void loop() {
  if (irqFlag) {
    noInterrupts();
    irqFlag = false;
    interrupts();

    handleLightningInterrupt();
  }

  resetTransientStatusIfTimedOut();
  resetStrikeDisplayIfTimedOut();

  if (millis() - lastRefreshMillis >= 1000UL) {
    lastRefreshMillis = millis();
    drawValues();
  }

  // Line-powered application: backlight remains continuously ON.
}