/**
 * RemoteRig — ESP32 MQTT Bridge Firmware
 * ======================================
 * Dedicated board per camera node. Connects the ESP8266 camera bridge
 * to the RemoteRig MQTT hub.
 *
 * ONE JOB: relay between UART (ESP8266) and MQTT (Pi hub).
 *   - Connects to travel router Wi-Fi
 *   - Reads status JSON from ESP8266 over UART → publishes via MQTT
 *   - Receives commands via MQTT from hub → forwards to ESP8266 over UART
 *   - Handles auto-registration (announce on first boot)
 *   - Heartbeat publishing
 *   - Zero camera communication, zero network switching
 *
 * UART Protocol: JSON-per-line at 115200 8N1
 *   ESP8266 → ESP32:  {"type":"status","battery_raw":217,...}\n
 *   ESP8266 → ESP32:  {"type":"ack","cmd":"start_recording"}\n
 *   ESP32 → ESP8266:  {"type":"cmd","command":"start_recording"}\n
 *   ESP32 → ESP8266:  {"type":"cmd","command":"ping"}\n
 *
 * Hardware:
 *   - Seeed Studio XIAO ESP32-C6
 *   - Serial1: RX=D7, TX=D6 (crossed to the ESP-01S TX/RX)
 *   - Shared GND between boards
 *   - 5V rail → XIAO 5V/VIN; ESP-01S on its own 3.3V buck
 */

#include <Arduino.h>
#include <WiFi.h>
#include <WiFiClient.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <LittleFS.h>
#include <Wire.h>
#include <U8g2lib.h>

// ────────────────────────────────────────────
// Configuration (LittleFS)
// ────────────────────────────────────────────

struct Config {
  String wifi_ssid       = "RemoteRig";
  String wifi_password   = "";
  String mqtt_broker     = "192.168.8.56";
  int    mqtt_port       = 1883;
  String camera_id       = "";   // assigned by hub
  int    heartbeat_sec   = 60;
  // Battery calibration: two-point linear map of the GoPro offset-57
  // raw byte → percent. Uncalibrated when max <= min (then we omit
  // battery_pct per the MQTT contract). Set via the set_battery_cal
  // command and persisted here.
  int    bat_raw_min     = 0;    // raw at 0%
  int    bat_raw_max     = 0;    // raw at 100%
} cfg;

bool loadConfig() {
  if (!LittleFS.begin(true)) { Serial.println("[CFG] LittleFS mount failed"); return false; }
  File f = LittleFS.open("/config.json", "r");
  if (!f) { Serial.println("[CFG] No config — using defaults"); return false; }

  JsonDocument doc;
  DeserializationError err = deserializeJson(doc, f);
  f.close();
  if (err) { Serial.printf("[CFG] Parse error: %s\n", err.c_str()); return false; }

  cfg.wifi_ssid       = doc["wifi_ssid"]     | cfg.wifi_ssid;
  cfg.wifi_password   = doc["wifi_password"] | cfg.wifi_password;
  cfg.mqtt_broker     = doc["mqtt_broker"]   | cfg.mqtt_broker;
  cfg.mqtt_port       = doc["mqtt_port"]     | cfg.mqtt_port;
  cfg.camera_id       = doc["camera_id"]     | cfg.camera_id;
  cfg.heartbeat_sec   = doc["heartbeat_interval_sec"] | cfg.heartbeat_sec;
  cfg.bat_raw_min     = doc["bat_raw_min"]   | cfg.bat_raw_min;
  cfg.bat_raw_max     = doc["bat_raw_max"]   | cfg.bat_raw_max;
  Serial.printf("[CFG] Loaded: ssid=%s broker=%s:%d cam=%s\n",
                cfg.wifi_ssid.c_str(), cfg.mqtt_broker.c_str(), cfg.mqtt_port,
                cfg.camera_id.length() ? cfg.camera_id.c_str() : "-");
  return true;
}

bool saveConfig() {
  File f = LittleFS.open("/config.json", "w");
  if (!f) return false;
  JsonDocument doc;
  doc["wifi_ssid"]     = cfg.wifi_ssid;
  doc["wifi_password"] = cfg.wifi_password;
  doc["mqtt_broker"]   = cfg.mqtt_broker;
  doc["mqtt_port"]     = cfg.mqtt_port;
  doc["camera_id"]     = cfg.camera_id;
  doc["heartbeat_interval_sec"] = cfg.heartbeat_sec;
  doc["bat_raw_min"]   = cfg.bat_raw_min;
  doc["bat_raw_max"]   = cfg.bat_raw_max;
  serializeJson(doc, f);
  f.close();
  return true;
}

// Map a raw offset-57 byte to battery percent using the stored
// two-point calibration. Returns -1 when uncalibrated.
int batteryPct(int raw) {
  if (cfg.bat_raw_max <= cfg.bat_raw_min) return -1;   // uncalibrated
  long pct = (long)(raw - cfg.bat_raw_min) * 100 /
             (cfg.bat_raw_max - cfg.bat_raw_min);
  if (pct < 0)   pct = 0;
  if (pct > 100) pct = 100;
  return (int)pct;
}

// ────────────────────────────────────────────
// UART to ESP-01S (HardwareSerial1)
// ────────────────────────────────────────────
// XIAO ESP32-C6 Serial1: RX=D7, TX=D6 (Serial = native USB CDC)
// Connect: XIAO RX(D7) ← ESP-01S TX
//          XIAO TX(D6) → ESP-01S RX

#define UART_ESP8266 Serial1
#define UART_RX_PIN  D7
#define UART_TX_PIN  D6

// Reserved for future ESP-01S UART OTA ("XIAO as flasher"): two control
// lines let the XIAO drive the ESP-01S into its serial bootloader and
// reflash it over Serial1 — no USB-UART adapter or GPIO0 jumper needed.
// Not driven yet; see docs/design/esp01s-uart-ota.md.
#define ESP01_RST_PIN D8    // → ESP-01S RST   (pulse low to reset)
#define ESP01_PGM_PIN D10   // → ESP-01S GPIO0 (low at reset = bootloader)

// ────────────────────────────────────────────
// RGB STAT LED — D0/D1/D2 (red/green/blue) via 220Ω each
// ────────────────────────────────────────────
// Wiring assumes common cathode (HIGH = on). Set RGB_COMMON_ANODE to
// 1 for a common-anode part (LOW = on).
#define RGB_PIN_R        D0
#define RGB_PIN_G        D1
#define RGB_PIN_B        D2
#define RGB_COMMON_ANODE 1   // this module is common-anode (LOW = on)

void rgbWrite(bool r, bool g, bool b) {
#if RGB_COMMON_ANODE
  digitalWrite(RGB_PIN_R, !r); digitalWrite(RGB_PIN_G, !g); digitalWrite(RGB_PIN_B, !b);
#else
  digitalWrite(RGB_PIN_R, r);  digitalWrite(RGB_PIN_G, g);  digitalWrite(RGB_PIN_B, b);
#endif
}

void rgbInit() {
  pinMode(RGB_PIN_R, OUTPUT);
  pinMode(RGB_PIN_G, OUTPUT);
  pinMode(RGB_PIN_B, OUTPUT);
  rgbWrite(0, 0, 1); // boot = blue
}

// ────────────────────────────────────────────
// Status OLED — 1.3" I2C panel on D4(SDA)/D5(SCL)
// ────────────────────────────────────────────
// 1.3" 128x64 modules are SH1106. If the image is shifted ~2px or
// wrapped, the panel is an SSD1306 — swap the constructor below to
// U8G2_SSD1306_128X64_NONAME_F_HW_I2C.
#define OLED_SDA_PIN  D4
#define OLED_SCL_PIN  D5
#define OLED_I2C_ADDR 0x3C

U8G2_SH1106_128X64_NONAME_F_HW_I2C oled(U8G2_R0, U8X8_PIN_NONE);
bool oledReady = false;

// Last-known camera status, mirrored for the display.
int  dispBatteryRaw  = 0;
bool dispRecording   = false;
int  dispVideoRemain = 0;          // seconds
unsigned long recStartMs = 0;      // 0 = not recording

// Walk the bus and log every responder — confirms the OLED address
// (and wiring) independent of the display driver.
void i2cScan() {
  Serial.println("[I2C] Scanning...");
  byte found = 0;
  for (byte a = 1; a < 127; a++) {
    Wire.beginTransmission(a);
    if (Wire.endTransmission() == 0) {
      Serial.printf("[I2C]   device @ 0x%02X\n", a);
      found++;
    }
  }
  if (!found) Serial.println("[I2C]   none found — check wiring/power");
}

void displayInit() {
  Wire.begin(OLED_SDA_PIN, OLED_SCL_PIN);
  i2cScan();
  oled.setI2CAddress(OLED_I2C_ADDR << 1);
  oledReady = oled.begin();
  Serial.printf("[OLED] begin %s\n", oledReady ? "ok" : "FAILED");
  if (!oledReady) return;
  oled.clearBuffer();
  oled.setFont(u8g2_font_7x14B_tr);
  oled.drawStr(0, 14, "RemoteRig");
  oled.setFont(u8g2_font_6x10_tr);
  oled.drawStr(0, 32, "Camera node");
  oled.drawStr(0, 46, "booting...");
  oled.sendBuffer();
}

void sendCmdToESP8266(const String& command) {
  JsonDocument doc;
  doc["type"]    = "cmd";
  doc["command"] = command;
  String line;
  serializeJson(doc, line);
  UART_ESP8266.println(line);
  UART_ESP8266.flush();
}

String uartLine;
bool readFromESP8266(String& line) {
  while (UART_ESP8266.available()) {
    char c = UART_ESP8266.read();
    if (c == '\n') {
      line = uartLine;
      uartLine = "";
      return true;
    }
    if (c != '\r') uartLine += c;
  }
  return false;
}

// ────────────────────────────────────────────
// MQTT
// ────────────────────────────────────────────

WiFiClient   routerClient;
PubSubClient mqtt(routerClient);

unsigned long bootMs = 0;
bool cameraOnline    = false;
unsigned long lastStatusMs = 0;

String clientID() {
  uint8_t mac[6];
  WiFi.macAddress(mac);
  char buf[32];
  snprintf(buf, sizeof(buf), "rig-%02x%02x%02x", mac[3], mac[4], mac[5]);
  return String(buf);
}

String mqttTopic(const char* t) {
  return "remoterig/cameras/" + cfg.camera_id + "/" + t;
}

void mqttCallback(char* topic, byte* payload, unsigned int len) {
  char buf[256];
  unsigned int n = len < 255 ? len : 255;
  memcpy(buf, payload, n); buf[n] = 0;

  JsonDocument doc;
  if (deserializeJson(doc, buf)) return;

  String cmd = doc["command"] | "";
  if (cmd == "start_recording" || cmd == "stop_recording") {
    Serial.printf("[MQTT] Forwarding command: %s → ESP8266\n", cmd.c_str());
    sendCmdToESP8266(cmd);
  } else if (cmd == "reboot") {
    ESP.restart();
  } else if (cmd == "set_battery_cal") {
    // Two ways to calibrate:
    //   explicit:  {"raw_min":185,"raw_max":245}
    //   capture:   {"point":"full"|"empty"}  → uses the latest raw reading
    String point = doc["point"] | "";
    if (point == "full")       cfg.bat_raw_max = dispBatteryRaw;
    else if (point == "empty") cfg.bat_raw_min = dispBatteryRaw;
    else {
      cfg.bat_raw_min = doc["raw_min"] | cfg.bat_raw_min;
      cfg.bat_raw_max = doc["raw_max"] | cfg.bat_raw_max;
    }
    saveConfig();
    Serial.printf("[BAT] Calibration set: raw_min=%d raw_max=%d\n",
                  cfg.bat_raw_min, cfg.bat_raw_max);
  } else if (cmd == "set_camera_config") {
    // Forward camera-bridge config to the ESP-01S over UART so the
    // GoPro creds / poll rate can change without reflashing it.
    JsonDocument out;
    out["type"]    = "cmd";
    out["command"] = "set_config";
    if (!doc["camera_ssid"].isNull())       out["camera_ssid"]       = doc["camera_ssid"];
    if (!doc["camera_password"].isNull())   out["camera_password"]   = doc["camera_password"];
    if (!doc["camera_ip"].isNull())         out["camera_ip"]         = doc["camera_ip"];
    if (!doc["poll_interval_sec"].isNull()) out["poll_interval_sec"] = doc["poll_interval_sec"];
    String line; serializeJson(out, line);
    UART_ESP8266.println(line);
    UART_ESP8266.flush();
    Serial.println("[MQTT] Forwarded set_config → ESP-01S");
  } else if (cmd == "registered") {
    String id = doc["camera_id"] | "";
    if (id.length() > 0 && id != cfg.camera_id) {
      cfg.camera_id = id;
      saveConfig();
      mqtt.unsubscribe(mqttTopic("command").c_str());
      mqtt.subscribe(mqttTopic("command").c_str(), 2);
      Serial.printf("[MQTT] Registered as %s\n", id.c_str());
    }
  }
}

bool connectMQTT() {
  mqtt.setServer(cfg.mqtt_broker.c_str(), cfg.mqtt_port);
  mqtt.setCallback(mqttCallback);
  mqtt.setKeepAlive(60);

  if (!mqtt.connect(clientID().c_str())) {
    Serial.printf("[MQTT] Connect fail (state=%d)\n", mqtt.state());
    return false;
  }

  Serial.println("[MQTT] Connected");

  // Subscribe to commands (if registered)
  if (cfg.camera_id.length() > 0) {
    mqtt.subscribe(mqttTopic("command").c_str(), 2);
  }

  // Announce if new
  if (cfg.camera_id.length() == 0) {
    JsonDocument doc;
    doc["mac_address"]      = WiFi.macAddress();
    doc["firmware_version"] = "0.3.0-esp32-mqtt-bridge";
    doc["friendly_name"]    = "Cam-" + clientID();
    JsonArray caps = doc["capabilities"].to<JsonArray>();
    caps.add("start_stop"); caps.add("status");
    String payload; serializeJson(doc, payload);
    String announceTopic = "remoterig/cameras/announce-" + clientID();
    mqtt.publish(announceTopic.c_str(), payload.c_str(), true);
    Serial.println("[MQTT] Announced for registration");
  }

  return true;
}

// ────────────────────────────────────────────
// Status screen + LED
// ────────────────────────────────────────────

// Reflect overall health on the RGB STAT LED.
//   red     = offline (no Wi-Fi)
//   magenta = Wi-Fi up, hub (MQTT) unreachable
//   yellow  = hub up, GoPro unreachable
//   green   = healthy (hub + camera reachable)
void updateStatusLed() {
  if (WiFi.status() != WL_CONNECTED) rgbWrite(1, 0, 0);       // red
  else if (!mqtt.connected())        rgbWrite(1, 0, 1);       // magenta
  else if (!cameraOnline)            rgbWrite(1, 1, 0);       // yellow
  else                               rgbWrite(0, 1, 0);       // green
}

void renderStatus() {
  if (!oledReady) return;
  oled.clearBuffer();

  // Camera id (top, bold)
  oled.setFont(u8g2_font_7x14B_tr);
  String id = cfg.camera_id.length() ? cfg.camera_id : clientID();
  oled.drawStr(0, 13, id.c_str());

  oled.setFont(u8g2_font_6x10_tr);
  char line[24];

  // REC state + session timer
  if (dispRecording) {
    unsigned long s = recStartMs ? (millis() - recStartMs) / 1000 : 0;
    oled.drawBox(0, 19, 6, 6);                       // filled square = REC
    snprintf(line, sizeof(line), "REC  %02lu:%02lu", s / 60, s % 60);
    oled.drawStr(10, 26, line);
  } else {
    oled.drawStr(0, 26, "IDLE");
  }

  // Battery (% when calibrated, else raw) + video remaining (minutes)
  int pct = batteryPct(dispBatteryRaw);
  if (pct >= 0) snprintf(line, sizeof(line), "BAT %d%%  VID %dm", pct, dispVideoRemain / 60);
  else          snprintf(line, sizeof(line), "BAT %d  VID %dm", dispBatteryRaw, dispVideoRemain / 60);
  oled.drawStr(0, 38, line);

  // Uplink to the hub
  const char* link = mqtt.connected()              ? "LINK: MQTT ok"
                   : WiFi.status() == WL_CONNECTED  ? "LINK: wifi only"
                                                    : "LINK: offline";
  oled.drawStr(0, 50, link);

  // Camera reachability
  oled.drawStr(0, 62, cameraOnline ? "CAM: online" : "CAM: --");

  oled.sendBuffer();
}

// ────────────────────────────────────────────
// Setup
// ────────────────────────────────────────────

void setup() {
  Serial.begin(115200);
  delay(500);
  Serial.println("\n[BRIDGE] ESP32 MQTT Bridge v1.0");

  bootMs = millis();
  rgbInit();     // RGB STAT LED — blue during boot

  displayInit(); // I2C scan + OLED splash

  loadConfig();

  // UART to ESP-01S
  UART_ESP8266.begin(115200, SERIAL_8N1, UART_RX_PIN, UART_TX_PIN);
  Serial.println("[UART] ESP-01S link on Serial1 (RX=D7, TX=D6) @ 115200");

  // Connect to travel router — the ONLY network we touch
  Serial.printf("[WIFI] Connecting to: %s\n", cfg.wifi_ssid.c_str());
  WiFi.mode(WIFI_STA);
  WiFi.begin(cfg.wifi_ssid.c_str(), cfg.wifi_password.c_str());

  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 40) {
    delay(500); Serial.print("."); attempts++;
  }

  if (WiFi.status() == WL_CONNECTED) {
    Serial.printf("\n[WIFI] Connected. IP: %s\n", WiFi.localIP().toString().c_str());
  } else {
    Serial.println("\n[WIFI] FAILED — will retry");
  }

  // MQTT
  if (WiFi.status() == WL_CONNECTED) {
    connectMQTT();
  }
}

// ────────────────────────────────────────────
// Main Loop
// ────────────────────────────────────────────

void loop() {
  unsigned long now = millis();
  static unsigned long lastBeat = 0, lastRecon = 0;
  static int reconDelay = 1;

  // ── OLED + LED refresh (always — keep them live even when offline) ──
  static unsigned long lastDisp = 0;
  if (now - lastDisp > 500) { lastDisp = now; renderStatus(); updateStatusLed(); }

  // ── Wi-Fi watchdog ──
  if (WiFi.status() != WL_CONNECTED) {
    if (now - lastRecon > 5000) { lastRecon = now; WiFi.reconnect(); }
    delay(100); return;
  }

  // ── MQTT watchdog ──
  if (!mqtt.connected()) {
    if (now - lastRecon > (unsigned long)(reconDelay * 1000)) {
      lastRecon = now;
      if (connectMQTT()) reconDelay = 1;
      else reconDelay = min(reconDelay * 2, 30);
    }
    mqtt.loop(); delay(100); return;
  }

  mqtt.loop();

  // ── Read status from ESP8266 over UART → publish via MQTT ──
  String line;
  while (readFromESP8266(line)) {
    JsonDocument doc;
    DeserializationError err = deserializeJson(doc, line);
    if (err) { Serial.printf("[UART] Bad JSON: %s\n", line.c_str()); continue; }

    String type = doc["type"] | "";
    
    if (type == "status") {
      // Relay camera status to MQTT hub
      lastStatusMs = now;
      bool online = doc["online"] | false;
      cameraOnline = online;   // reflected on the RGB LED by updateStatusLed()

      // Mirror status onto the OLED fields
      dispBatteryRaw  = doc["battery_raw"] | 0;
      dispVideoRemain = doc["video_remaining_sec"] | 0;
      bool rec = doc["recording"] | false;
      if (rec && !dispRecording) recStartMs = millis();
      if (!rec)                  recStartMs = 0;
      dispRecording = rec;

      if (cfg.camera_id.length() > 0) {
        // Build the MQTT status payload per contract
        JsonDocument mqttDoc;
        mqttDoc["camera_id"]  = cfg.camera_id;
        mqttDoc["timestamp"]  = millis();
        mqttDoc["battery_raw"] = dispBatteryRaw;
        int pct = batteryPct(dispBatteryRaw);
        if (pct >= 0) mqttDoc["battery_pct"] = pct;   // omit when uncalibrated
        mqttDoc["video_remaining_sec"] = doc["video_remaining_sec"] | 0;
        mqttDoc["recording"]  = doc["recording"] | false;
        mqttDoc["online"]     = online;
        
        String payload;
        serializeJson(mqttDoc, payload);
        mqtt.publish(mqttTopic("status").c_str(), payload.c_str(), true);
      }
    }
    else if (type == "ack") {
      Serial.printf("[UART] ESP8266 ack: %s\n", doc["cmd"] | "?");
    }
    else if (type == "pong") {
      Serial.printf("[UART] ESP8266 pong (uptime=%d)\n", doc["uptime_ms"] | 0);
    }
    else if (type == "error") {
      Serial.printf("[UART] ESP8266 error: %s\n", doc["msg"] | "?");
    }
  }

  // ── Heartbeat to hub (every heartbeat_sec) ──
  if (cfg.camera_id.length() > 0 &&
      now - lastBeat > (unsigned long)(cfg.heartbeat_sec * 1000)) {
    lastBeat = now;
    JsonDocument doc;
    doc["camera_id"]  = cfg.camera_id;
    doc["timestamp"]  = millis();
    doc["uptime_sec"] = (now - bootMs) / 1000;
    doc["free_heap"]  = ESP.getFreeHeap();
    doc["status_age_ms"] = now - lastStatusMs;
    String payload; serializeJson(doc, payload);
    mqtt.publish(mqttTopic("heartbeat").c_str(), payload.c_str(), false);
  }

  // ── Periodic ping to ESP8266 to verify UART link ──
  static unsigned long lastPing = 0;
  if (now - lastPing > 30000) {
    lastPing = now;
    sendCmdToESP8266("ping");
  }

  delay(50);
}
