/*
 * Water for Aduvan — Water Purification Controller Firmware
 * Board: Seeed Studio XIAO ESP32-C3
 *
 * Responsibilities:
 *  - Read turbidity + flow rate sensors
 *  - Automatic control of pump + solenoid valve based on turbidity threshold
 *  - Battery voltage/percentage monitoring
 *  - Publish telemetry + heartbeat status over MQTT
 *  - Receive remote commands (pump, valve, mode, threshold, restart, OTA)
 *  - Wi-Fi + MQTT auto-reconnect
 *  - Local buffering of readings when offline, flushed on reconnect
 *  - OTA-ready structure (stub call-out, see NOTE below)
 *
 * Required libraries (Arduino Library Manager):
 *  - WiFi (bundled with ESP32 core)
 *  - PubSubClient by Nick O'Leary
 *  - ArduinoJson by Benoit Blanchon
 *  - Preferences (bundled, for local persistence of counters)
 *
 * Board package: install "esp32" by Espressif in Boards Manager,
 * select board "XIAO_ESP32C3".
 */

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

// ================== USER CONFIG ==================
const char* WIFI_SSID     = "WuodAwuor";
const char* WIFI_PASSWORD = "0748983442";

const char* MQTT_HOST = "10.30.74.227";
const int MQTT_PORT = 1883;
const char* MQTT_USER     = "";                // leave blank if allow_anonymous
const char* MQTT_PASS     = "";

const char* DEVICE_ID     = "WP-DEV-001";      // must match device_uid registered in the dashboard
const char* FIRMWARE_VERSION = "1.0.0";
// ===================================================

// ---- Pin mapping (XIAO ESP32-C3) ----
#define PIN_TURBIDITY   A2   // analog turbidity sensor (e.g. SEN0189)
#define PIN_FLOW        A1   // flow sensor pulse output (interrupt capable)
#define PIN_BATTERY_ADC D10   // battery voltage divider input
#define PIN_PUMP_RELAY  D9  // relay/MOSFET driving the 12V pump
#define PIN_VALVE_RELAY D8  // relay/MOSFET driving the 12V solenoid valve
#define PIN_EXT_POWER   D4   // digital input: HIGH when external/solar power present

// EXTENSION (future expansion): water level sensor input pin, LoRa module SPI/UART pins.
// Keep sensor-read and MQTT-publish functions separate (as below) so a LoRa transport
// can be swapped in without touching control logic.

WiFiClient wifiClient;
PubSubClient mqtt(wifiClient);
Preferences prefs;

// ---- Runtime state ----
bool pumpOn = false;
bool valveOpen = false;
bool autoMode = true;
float turbidityThresholdNTU = 5.0;

volatile unsigned long flowPulseCount = 0;
float currentFlowLpm = 0;
double litersToday = 0;
double litersMonth = 0;
double litersLifetime = 0;

unsigned long lastTelemetryMs = 0;
unsigned long lastHeartbeatMs = 0;
unsigned long lastFlowCalcMs = 0;
const unsigned long TELEMETRY_INTERVAL_MS = 10000;   // 10s sensor publish
const unsigned long HEARTBEAT_INTERVAL_MS = 30000;   // 30s heartbeat (within 30-60s spec)

// Flow sensor calibration: pulses per liter (datasheet-specific, e.g. YF-S201 ~450 P/L)
const float PULSES_PER_LITER = 450.0;

// Local offline buffer (simple ring buffer of JSON strings)
#define BUFFER_SIZE 50
String offlineBuffer[BUFFER_SIZE];
int bufferHead = 0;
int bufferCount = 0;

void IRAM_ATTR onFlowPulse() {
  flowPulseCount++;
}

// ---------------- Setup ----------------
void setup() {
  Serial.begin(115200);
  pinMode(PIN_PUMP_RELAY, OUTPUT);
  pinMode(PIN_VALVE_RELAY, OUTPUT);
  pinMode(PIN_EXT_POWER, INPUT);
  pinMode(PIN_FLOW, INPUT_PULLUP);
  attachInterrupt(digitalPinToInterrupt(PIN_FLOW), onFlowPulse, RISING);

  digitalWrite(PIN_PUMP_RELAY, LOW);
  digitalWrite(PIN_VALVE_RELAY, LOW);

  prefs.begin("aquactrl", false);
  litersLifetime = prefs.getDouble("liters_life", 0);
  litersMonth = prefs.getDouble("liters_month", 0);
  turbidityThresholdNTU = prefs.getFloat("threshold", 5.0);

  connectWiFi();
  mqtt.setServer(MQTT_HOST, MQTT_PORT);
  mqtt.setBufferSize(512);      // <-- NEW: raises MQTT packet limit above the 256-byte default
  mqtt.setCallback(onMqttMessage);
  connectMqtt();

  lastFlowCalcMs = millis();
  Serial.println("Water for Aduvan firmware ready.");
}
// ---------------- Main loop ----------------
void loop() {
  if (WiFi.status() != WL_CONNECTED) {
    connectWiFi();
  }
  if (!mqtt.connected()) {
    connectMqtt();
  } else {
    mqtt.loop();
    flushOfflineBuffer();
  }

  calculateFlow();
  runControlLogic();

  unsigned long now = millis();
  if (now - lastTelemetryMs >= TELEMETRY_INTERVAL_MS) {
    lastTelemetryMs = now;
    publishTelemetry();
  }
  if (now - lastHeartbeatMs >= HEARTBEAT_INTERVAL_MS) {
    lastHeartbeatMs = now;
    publishStatus();
  }
}

// ---------------- Wi-Fi / MQTT connectivity ----------------
void connectWiFi() {
  if (WiFi.status() == WL_CONNECTED) return;
  Serial.println("Connecting to Wi-Fi...");
  WiFi.mode(WIFI_STA);
  WiFi.begin(WIFI_SSID, WIFI_PASSWORD);
  unsigned long start = millis();
  while (WiFi.status() != WL_CONNECTED && millis() - start < 15000) {
    delay(300);
  }
  if (WiFi.status() == WL_CONNECTED) {
    Serial.println("Wi-Fi connected: " + WiFi.localIP().toString());
  } else {
    Serial.println("Wi-Fi connection failed, will retry.");
  }
}

void connectMqtt() {
  if (WiFi.status() != WL_CONNECTED) return;
  if (mqtt.connected()) return;
  String clientId = String("xiao-") + DEVICE_ID;
  Serial.println("Connecting to MQTT...");
  bool ok = (strlen(MQTT_USER) > 0)
    ? mqtt.connect(clientId.c_str(), MQTT_USER, MQTT_PASS)
    : mqtt.connect(clientId.c_str());
  if (ok) {
    Serial.println("MQTT connected.");
    String cmdTopic = String("devices/") + DEVICE_ID + "/commands";
    mqtt.subscribe(cmdTopic.c_str());
  } else {
    Serial.println("MQTT connect failed, rc=" + String(mqtt.state()));
  }
}

// ---------------- Sensor reading ----------------
float readTurbidityNTU() {
  int raw = analogRead(PIN_TURBIDITY);
  // EXTENSION: replace with your sensor's calibration curve (raw ADC -> NTU).
  float voltage = raw * (3.3 / 4095.0);
  float ntu = (2.5 - voltage) * 20.0; // placeholder linear approximation
  return ntu < 0 ? 0 : ntu;
}

float readBatteryVoltage() {
  int raw = analogRead(PIN_BATTERY_ADC);
  // EXTENSION: adjust divider ratio to match your battery monitoring circuit.
  float dividerRatio = 2.0;
  return raw * (3.3 / 4095.0) * dividerRatio;
}

float batteryPercentFromVoltage(float v) {
  // Rough LiPo/lead-acid curve placeholder — replace with your battery's real curve.
  float minV = 10.5, maxV = 12.6;
  float pct = (v - minV) / (maxV - minV) * 100.0;
  if (pct < 0) pct = 0;
  if (pct > 100) pct = 100;
  return pct;
}

void calculateFlow() {
  unsigned long now = millis();
  unsigned long elapsed = now - lastFlowCalcMs;
  if (elapsed < 1000) return; // update at most once per second

  noInterrupts();
  unsigned long pulses = flowPulseCount;
  flowPulseCount = 0;
  interrupts();

  float liters = pulses / PULSES_PER_LITER;
  currentFlowLpm = (liters / (elapsed / 1000.0)) * 60.0;

  litersToday += liters;
  litersMonth += liters;
  litersLifetime += liters;

  lastFlowCalcMs = now;

  // Persist periodically (every ~1 min via heartbeat cadence would reduce flash wear;
  // simplified here to persist on every calc for clarity).
  prefs.putDouble("liters_life", litersLifetime);
  prefs.putDouble("liters_month", litersMonth);
}

// ---------------- Automatic control logic ----------------
void runControlLogic() {
  if (!autoMode) return; // manual mode: controls only driven by MQTT commands

  float turbidity = readTurbidityNTU();

  if (turbidity > turbidityThresholdNTU) {
    // Still dirty: keep filtering, keep output valve closed
    setPump(true);
    setValve(false);
  } else {
    // Clean enough: allow output
    setValve(true);
  }

  // No-flow-while-pump-on fault is raised backend-side from telemetry,
  // but we also stop the pump locally as a safety fallback after a timeout.
  static unsigned long noFlowSinceMs = 0;
  if (pumpOn && currentFlowLpm <= 0.05) {
    if (noFlowSinceMs == 0) noFlowSinceMs = millis();
    if (millis() - noFlowSinceMs > 30000) { // 30s of no flow while pump is on
      setPump(false);
      publishEvent("pump_failure", "Pump stopped: no flow detected for 30s");
      noFlowSinceMs = 0;
    }
  } else {
    noFlowSinceMs = 0;
  }
}

void setPump(bool on) {
  pumpOn = on;
  digitalWrite(PIN_PUMP_RELAY, on ? HIGH : LOW);
}

void setValve(bool open) {
  valveOpen = open;
  digitalWrite(PIN_VALVE_RELAY, open ? HIGH : LOW);
}

// ---------------- MQTT publish ----------------
void publishTelemetry() {
  StaticJsonDocument<384> doc;
  doc["turbidity_ntu"] = readTurbidityNTU();
  doc["flow_rate_lpm"] = currentFlowLpm;
  doc["liters_today"] = litersToday;
  doc["liters_month"] = litersMonth;
  doc["liters_lifetime"] = litersLifetime;
  float vbat = readBatteryVoltage();
  doc["battery_v"] = vbat;
  doc["battery_pct"] = batteryPercentFromVoltage(vbat);
  doc["external_power"] = digitalRead(PIN_EXT_POWER) == HIGH;
  doc["wifi_rssi"] = WiFi.RSSI();
  doc["pump_on"] = pumpOn;
  doc["valve_open"] = valveOpen;
  doc["firmware_version"] = FIRMWARE_VERSION;

  String payload;
  serializeJson(doc, payload);
  String topic = String("devices/") + DEVICE_ID + "/telemetry";

if (mqtt.connected()) {
    bool ok = mqtt.publish(topic.c_str(), payload.c_str());
    if (!ok) Serial.println("Telemetry publish FAILED (packet too large?)");
  } else {
    bufferOffline(payload);
  }
}

void publishStatus() {
  StaticJsonDocument<256> doc;
  doc["uptime_s"] = millis() / 1000;
  doc["firmware_version"] = FIRMWARE_VERSION;
  doc["mode"] = autoMode ? "auto" : "manual";
  doc["wifi_rssi"] = WiFi.RSSI();
  doc["mqtt_connected"] = mqtt.connected();

  String payload;
  serializeJson(doc, payload);
  String topic = String("devices/") + DEVICE_ID + "/status";
  if (mqtt.connected()) mqtt.publish(topic.c_str(), payload.c_str());
}

void publishEvent(const char* type, const char* message) {
  StaticJsonDocument<256> doc;
  doc["type"] = type;
  doc["message"] = message;
  String payload;
  serializeJson(doc, payload);
  String topic = String("devices/") + DEVICE_ID + "/status"; // reuse status channel for events
  if (mqtt.connected()) mqtt.publish(topic.c_str(), payload.c_str());
}

// ---------------- Offline buffering ----------------
void bufferOffline(const String& payload) {
  offlineBuffer[bufferHead] = payload;
  bufferHead = (bufferHead + 1) % BUFFER_SIZE;
  if (bufferCount < BUFFER_SIZE) bufferCount++;
}

void flushOfflineBuffer() {
  if (bufferCount == 0) return;
  String topic = String("devices/") + DEVICE_ID + "/telemetry";
  int start = (bufferHead - bufferCount + BUFFER_SIZE) % BUFFER_SIZE;
  for (int i = 0; i < bufferCount; i++) {
    int idx = (start + i) % BUFFER_SIZE;
    mqtt.publish(topic.c_str(), offlineBuffer[idx].c_str());
  }
  bufferCount = 0;
  bufferHead = 0;
}

// ---------------- Command handling ----------------
void onMqttMessage(char* topic, byte* payload, unsigned int length) {
  String msg;
  for (unsigned int i = 0; i < length; i++) msg += (char)payload[i];

  StaticJsonDocument<256> doc;
  DeserializationError err = deserializeJson(doc, msg);
  if (err) return;

  String command = doc["command"] | "";
  String commandId = doc["command_id"] | "";

  bool success = true;

  if (command == "pump_on") setPump(true);
  else if (command == "pump_off") setPump(false);
  else if (command == "valve_open") setValve(true);
  else if (command == "valve_close") setValve(false);
  else if (command == "set_mode") autoMode = (String((const char*)(doc["mode"] | "auto")) == "auto");
  else if (command == "set_threshold") {
    turbidityThresholdNTU = doc["turbidity_threshold"] | turbidityThresholdNTU;
    prefs.putFloat("threshold", turbidityThresholdNTU);
  }
  else if (command == "restart") {
    ackCommand(commandId, true);
    delay(200);
    ESP.restart();
  }
  else if (command == "ota_update") {
    // NOTE (production): implement HTTPUpdate.update(wifiClient, file_url) here,
    // verifying the firmware checksum before flashing. Kept as a stub so this
    // sketch compiles without pulling in HTTPUpdate/HTTPClient unconditionally.
    publishEvent("ota_update", "OTA command received (not yet implemented in this build)");
    success = false;
  }
  else {
    success = false;
  }

  ackCommand(commandId, success);
}

void ackCommand(const String& commandId, bool success) {
  if (commandId.length() == 0) return;
  StaticJsonDocument<128> doc;
  doc["command_id"] = commandId;
  doc["success"] = success;
  String payload;
  serializeJson(doc, payload);
  String topic = String("devices/") + DEVICE_ID + "/commands/ack";
  if (mqtt.connected()) mqtt.publish(topic.c_str(), payload.c_str());
}