/*
 * ============================================================
 * Automated Chemical Dosing System - Hospital IoT
 * FILE: supply_node.ino
 * DESCRIPTION: Supply station controller
 *              - Monitors water tank level (ultrasonic sim)
 *              - Controls chemical pump (PWM)
 *              - Controls water solenoid valve
 *              - Executes dispensing cycles
 *              - Detects faults via flow monitoring
 *              - Displays water level on LCD
 * SIMULATION: WokWi ESP32
 * ============================================================
 */

#include <Wire.h>
#include <LiquidCrystal_I2C.h>

// ─── PIN DEFINITIONS ────────────────────────────────────────
#define PIN_WATER_VALVE     26    // Water solenoid valve (LED)
#define PIN_PUMP            14    // Chemical pump (LED + PWM)
#define PIN_ALARM_LED       13    // Red alarm indicator
#define PIN_BUZZER          12    // Audible alarm
#define PIN_FLOW_SIM        34    // Potentiometer simulates flow
#define PIN_FAULT_BTN       35    // Button simulates fault injection

// ─── SYSTEM CONSTANTS ───────────────────────────────────────
#define TANK_HEIGHT_M       10.18  // Calculated water tank height
#define TANK_RADIUS_M       0.25   // Water tank radius
#define TANK_MAX_VOL_L      2000.0 // Max water volume (litres)
#define CHEM_CONTAINER_ML   10000  // Chemical container (10L)
#define DISPENSE_TARGET_ML  600    // Target refill per cycle (mL)
#define TARGET_PPM          500    // Target concentration
#define MIN_WATER_LEVEL_PCT 10     // Low water warning threshold %
#define MIN_CHEM_ML         500    // Low chemical warning threshold
#define CYCLE_TIMEOUT_MS    30000  // Max dispensing cycle duration
#define PWM_CHANNEL         0      // ESP32 LEDC PWM channel
#define PWM_FREQ            1000   // PWM frequency (Hz)
#define PWM_RESOLUTION      8      // 8-bit resolution (0-255)
#define HEARTBEAT_INTERVAL  5000   // Heartbeat every 5s (sim, 30s real)
#define LCD_UPDATE_INTERVAL 1000   // LCD refresh interval

// ─── SYSTEM STATE ENUM ──────────────────────────────────────
enum SystemState {
  STATE_IDLE,
  STATE_DISPENSING,
  STATE_FAULT,
  STATE_LOW_SUPPLY
};

// ─── GLOBAL VARIABLES ───────────────────────────────────────
LiquidCrystal_I2C lcd(0x27, 16, 2);

SystemState currentState       = STATE_IDLE;
float       waterLevelPct      = 85.0;   // Simulated water level %
float       chemRemainingML    = 8500.0; // Simulated chemical remaining
float       waterFlowTotal     = 0.0;    // mL dispensed this cycle
float       chemFlowTotal      = 0.0;    // mL chemical this cycle
float       actualPPM          = 0.0;    // Calculated concentration
int         pumpPWM            = 128;    // Current pump PWM (0-255)
bool        cycleActive        = false;
bool        faultActive        = false;

unsigned long cycleStartTime   = 0;
unsigned long lastHeartbeat    = 0;
unsigned long lastLCDUpdate    = 0;
unsigned long lastSensorRead   = 0;

// Simulated MQTT queue (serial commands in simulation)
// In real system: WiFi + PubSubClient library
String pendingCommand          = "";

// ─── FUNCTION PROTOTYPES ────────────────────────────────────
void     initPeripherals();
void     updateLCD();
void     readSensors();
float    simulateWaterLevel();
float    readFlowMeter();
void     startDispensingCycle(int targetML);
void     runDispensingLoop();
void     stopCycle(bool success);
void     triggerFault(String reason);
void     clearFault();
void     checkSupplyLevels();
void     adjustPumpPWM();
void     publishHeartbeat();
void     publishCycleLog(bool success, float volDispensed);
void     checkSerialCommands();
void     activateAlarm(int priority);
void     deactivateAlarm();

// ────────────────────────────────────────────────────────────
void setup() {
  Serial.begin(115200);
  Serial.println(F("=== Chemical Dosing System - Supply Node ==="));
  Serial.println(F("Initializing..."));

  initPeripherals();

  // Simulate WiFi connection
  Serial.println(F("[MQTT] Connected to broker: raspberry-pi-central"));
  Serial.println(F("[MQTT] Subscribed: hospital/supply/command"));
  Serial.println(F("[SYSTEM] Supply node online. Entering IDLE state."));

  updateLCD();
}

// ────────────────────────────────────────────────────────────
void loop() {

  unsigned long now = millis();

  // ── Read serial commands (simulates MQTT) ──
  checkSerialCommands();

  // ── Read sensors periodically ──
  if (now - lastSensorRead >= 500) {
    lastSensorRead = now;
    readSensors();
    checkSupplyLevels();
  }

  // ── Update LCD periodically ──
  if (now - lastLCDUpdate >= LCD_UPDATE_INTERVAL) {
    lastLCDUpdate = now;
    updateLCD();
  }

  // ── Publish heartbeat periodically ──
  if (now - lastHeartbeat >= HEARTBEAT_INTERVAL) {
    lastHeartbeat = now;
    publishHeartbeat();
  }

  // ── Run dispensing loop if active ──
  if (currentState == STATE_DISPENSING) {
    runDispensingLoop();
  }

  // ── Check fault button (fault injection for testing) ──
  if (digitalRead(PIN_FAULT_BTN) == LOW) {
    delay(50); // debounce
    if (digitalRead(PIN_FAULT_BTN) == LOW) {
      triggerFault("MANUAL_FAULT_INJECTION");
    }
  }
}

// ────────────────────────────────────────────────────────────
void initPeripherals() {

  // Output pins
  pinMode(PIN_WATER_VALVE,  OUTPUT);
  pinMode(PIN_PUMP,         OUTPUT);
  pinMode(PIN_ALARM_LED,    OUTPUT);
  pinMode(PIN_BUZZER,       OUTPUT);

  // Input pins
  pinMode(PIN_FAULT_BTN,    INPUT_PULLUP);

  // Ensure all actuators start OFF (fail-safe)
  digitalWrite(PIN_WATER_VALVE, LOW);
  digitalWrite(PIN_PUMP,        LOW);
  digitalWrite(PIN_ALARM_LED,   LOW);
  digitalWrite(PIN_BUZZER,      LOW);

  // PWM setup for pump control
  ledcAttachChannel(PIN_PUMP, PWM_FREQ, PWM_RESOLUTION, PWM_CHANNEL);
  ledcWrite(PWM_CHANNEL, 0); // Start with pump off

  // LCD init
  lcd.init();
  lcd.backlight();
  lcd.setCursor(0, 0);
  lcd.print(F("Chem Dose System"));
  lcd.setCursor(0, 1);
  lcd.print(F("Initializing... "));
  delay(1500);
  lcd.clear();

  Serial.println(F("[INIT] Peripherals initialized."));
}

// ────────────────────────────────────────────────────────────
void readSensors() {
  /*
   * In real system:
   * - Read GPIO34 ADC for UB500 4-20mA → water level
   * - Read HX711 for load cell → chemical weight
   * - Count flow meter pulses via interrupts
   *
   * In simulation:
   * - Potentiometer on GPIO34 simulates flow rate
   * - Water level decremented during active cycles
   */

  // Simulate water level slowly decreasing over time
  if (currentState == STATE_DISPENSING) {
    waterLevelPct -= 0.1;
    if (waterLevelPct < 0) waterLevelPct = 0;
  }

  // Read potentiometer as simulated flow rate (0-100 mL/s)
  int adcVal = analogRead(PIN_FLOW_SIM);
  float simulatedFlowRate = (adcVal / 4095.0) * 100.0;

  // Accumulate flow totals during active cycle
  if (currentState == STATE_DISPENSING && cycleActive) {
    // Water flow increases faster than chemical (dilution ratio)
    waterFlowTotal += simulatedFlowRate * 0.9 * 0.5; // × 0.5s interval
    chemFlowTotal  += simulatedFlowRate * 0.1 * 0.5;
    chemRemainingML -= simulatedFlowRate * 0.1 * 0.5;

    // Calculate actual PPM
    if ((waterFlowTotal + chemFlowTotal) > 0) {
      actualPPM = (chemFlowTotal / (waterFlowTotal + chemFlowTotal)) * 1000000.0;
    }
  }
}

// ────────────────────────────────────────────────────────────
void checkSupplyLevels() {
  /*
   * Requirement 7: Warn when supply levels fall below threshold
   */

  // Water level warning
  if (waterLevelPct < MIN_WATER_LEVEL_PCT && currentState != STATE_FAULT) {
    Serial.println(F("[WARNING-P3] Water tank below 10% threshold!"));
    Serial.print(F("[MQTT] publish → hospital/alarms : "));
    Serial.println(F("{type:LOW_WATER, priority:P3}"));
    currentState = STATE_LOW_SUPPLY;
    activateAlarm(3);
  }

  // Chemical level warning
  if (chemRemainingML < MIN_CHEM_ML && currentState != STATE_FAULT) {
    Serial.println(F("[WARNING-P3] Chemical container below 500mL!"));
    Serial.print(F("[MQTT] publish → hospital/alarms : "));
    Serial.println(F("{type:LOW_CHEMICAL, priority:P3}"));
    activateAlarm(3);
  }
}

// ────────────────────────────────────────────────────────────
void startDispensingCycle(int targetML) {
  /*
   * Begin a dispensing cycle for a given target volume.
   * Opens water valve, starts pump, records start time.
   */

  if (faultActive) {
    Serial.println(F("[ERROR] Cannot start cycle - fault active. Clear fault first."));
    return;
  }

  if (waterLevelPct <= 0 || chemRemainingML <= 0) {
    triggerFault("DRY_SUPPLY_TANK");
    return;
  }

  Serial.println(F("─────────────────────────────────"));
  Serial.println(F("[CYCLE] Starting dispensing cycle"));
  Serial.print(F("[CYCLE] Target volume: "));
  Serial.print(targetML);
  Serial.println(F(" mL"));

  // Reset flow totals for new cycle
  waterFlowTotal = 0.0;
  chemFlowTotal  = 0.0;
  actualPPM      = 0.0;
  cycleStartTime = millis();
  cycleActive    = true;

  // Open water valve
  digitalWrite(PIN_WATER_VALVE, HIGH);
  Serial.println(F("[ACTUATOR] Water solenoid valve: OPEN"));

  // Start pump at initial PWM
  pumpPWM = 128; // 50% initial duty cycle
  ledcWrite(PWM_CHANNEL, pumpPWM);
  Serial.print(F("[ACTUATOR] Chemical pump: ON, PWM="));
  Serial.println(pumpPWM);

  currentState = STATE_DISPENSING;

  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print(F("DISPENSING...   "));
}

// ────────────────────────────────────────────────────────────
void runDispensingLoop() {
  /*
   * Called every loop iteration during active dispensing.
   * Checks:
   * 1. Target volume reached → stop cycle (success)
   * 2. Cycle timeout → fault (blocked valve / pump fail)
   * 3. Flow meter reads zero → pump/valve fault
   * 4. Adjusts pump PWM for concentration control
   */

  unsigned long elapsed = millis() - cycleStartTime;

  // ── Fault: cycle timeout ──
  if (elapsed > CYCLE_TIMEOUT_MS) {
    triggerFault("CYCLE_TIMEOUT - possible blocked valve or pump failure");
    return;
  }

  // ── Fault: no flow detected after 5 seconds ──
  if (elapsed > 5000 && waterFlowTotal < 1.0) {
    triggerFault("NO_FLOW_DETECTED - water valve or pump failure");
    return;
  }

  // ── Adjust pump PWM for target concentration ──
  adjustPumpPWM();

  // ── Check if target volume reached ──
  float totalDispensed = waterFlowTotal + chemFlowTotal;
  if (totalDispensed >= DISPENSE_TARGET_ML) {
    stopCycle(true);
    return;
  }

  // ── Serial progress update every 2 seconds ──
  static unsigned long lastPrint = 0;
  if (millis() - lastPrint >= 2000) {
    lastPrint = millis();
    Serial.print(F("[CYCLE] Dispensed: "));
    Serial.print(totalDispensed, 1);
    Serial.print(F(" / "));
    Serial.print(DISPENSE_TARGET_ML);
    Serial.print(F(" mL | PPM: "));
    Serial.print(actualPPM, 0);
    Serial.print(F(" | PWM: "));
    Serial.println(pumpPWM);

    // Update LCD with progress
    lcd.setCursor(0, 1);
    lcd.print((int)totalDispensed);
    lcd.print(F("/"));
    lcd.print(DISPENSE_TARGET_ML);
    lcd.print(F("mL "));
    lcd.print((int)actualPPM);
    lcd.print(F("ppm"));
  }
}

// ────────────────────────────────────────────────────────────
void adjustPumpPWM() {
  /*
   * Proportional concentration controller.
   * Adjusts pump PWM to maintain target PPM.
   * Error = target_ppm - actual_ppm
   * ΔP = Kp × error (clamped to ±20 per cycle)
   */

  if (waterFlowTotal < 10.0) return; // Need enough flow for stable reading

  float Kp    = 0.05;
  float error = TARGET_PPM - actualPPM;
  int   delta = (int)(Kp * error);

  // Clamp adjustment to prevent oscillation
  delta   = constrain(delta, -20, 20);
  pumpPWM = constrain(pumpPWM + delta, 30, 230);

  ledcWrite(PWM_CHANNEL, pumpPWM);
}

// ────────────────────────────────────────────────────────────
void stopCycle(bool success) {
  /*
   * Ends dispensing cycle.
   * Closes all actuators regardless of success/failure.
   */

  // Close all actuators (fail-safe: always close)
  digitalWrite(PIN_WATER_VALVE, LOW);
  ledcWrite(PWM_CHANNEL, 0);
  cycleActive = false;

  if (success) {
    Serial.println(F("─────────────────────────────────"));
    Serial.println(F("[CYCLE] Cycle COMPLETE - SUCCESS"));
    publishCycleLog(true, waterFlowTotal + chemFlowTotal);
    currentState = STATE_IDLE;
    deactivateAlarm();

    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print(F("Cycle Complete  "));
  } else {
    Serial.println(F("[CYCLE] Cycle ABORTED"));
    publishCycleLog(false, waterFlowTotal + chemFlowTotal);
    // State already set to FAULT by triggerFault()
  }

  Serial.println(F("[ACTUATOR] Water valve: CLOSED"));
  Serial.println(F("[ACTUATOR] Chemical pump: OFF"));
}

// ────────────────────────────────────────────────────────────
void triggerFault(String reason) {
  /*
   * Fault handler - called by any fault detection logic.
   * Stops all actuators, raises alarm, logs fault.
   */

  // Emergency stop all actuators
  digitalWrite(PIN_WATER_VALVE, LOW);
  ledcWrite(PWM_CHANNEL, 0);
  cycleActive  = false;
  faultActive  = true;
  currentState = STATE_FAULT;

  Serial.println(F("!!! FAULT DETECTED !!!"));
  Serial.print(F("[FAULT] Reason: "));
  Serial.println(reason);
  Serial.println(F("[MQTT] publish → hospital/alarms : {type:FAULT, priority:P1}"));

  // Determine alarm priority
  if (reason.indexOf("LEAK") >= 0 || reason.indexOf("DRY") >= 0) {
    activateAlarm(1); // P1 - Critical
  } else {
    activateAlarm(2); // P2 - High
  }

  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print(F("!! FAULT !!     "));
  lcd.setCursor(0, 1);
  lcd.print(reason.substring(0, 16));
}

// ────────────────────────────────────────────────────────────
void clearFault() {
  /*
   * Clears fault state after operator acknowledgement.
   * Requires serial command "CLEAR_FAULT"
   */
  faultActive  = false;
  currentState = STATE_IDLE;
  deactivateAlarm();
  Serial.println(F("[SYSTEM] Fault cleared by operator. System reset to IDLE."));
  updateLCD();
}

// ────────────────────────────────────────────────────────────
void activateAlarm(int priority) {
  /*
   * Priority 1 = continuous buzzer + red LED (critical)
   * Priority 2 = intermittent buzzer + red LED (high)
   * Priority 3 = red LED only, no buzzer (warning)
   */
  digitalWrite(PIN_ALARM_LED, HIGH);

  if (priority == 1) {
    digitalWrite(PIN_BUZZER, HIGH); // Continuous
    Serial.println(F("[ALARM] P1-CRITICAL: Continuous alarm activated"));
  } else if (priority == 2) {
    // Intermittent - handled in main loop (simplified here)
    digitalWrite(PIN_BUZZER, HIGH);
    Serial.println(F("[ALARM] P2-HIGH: Alarm activated"));
  } else {
    Serial.println(F("[ALARM] P3-WARNING: Indicator only"));
  }
}

// ────────────────────────────────────────────────────────────
void deactivateAlarm() {
  digitalWrite(PIN_ALARM_LED, LOW);
  digitalWrite(PIN_BUZZER,    LOW);
}

// ────────────────────────────────────────────────────────────
void updateLCD() {
  /*
   * Updates 16×2 LCD display with water level gauge.
   * Row 0: Water level bar + percentage
   * Row 1: System state + chemical remaining
   */

  lcd.setCursor(0, 0);
  lcd.print(F("W:"));

  // Draw 8-character bar graph
  int bars = (int)(waterLevelPct / 100.0 * 8);
  for (int i = 0; i < 8; i++) {
    lcd.print(i < bars ? '#' : '-');
  }

  lcd.print(F(" "));
  lcd.print((int)waterLevelPct);
  lcd.print(F("%  "));

  lcd.setCursor(0, 1);
  switch (currentState) {
    case STATE_IDLE:
      lcd.print(F("IDLE C:"));
      lcd.print((int)chemRemainingML);
      lcd.print(F("mL  "));
      break;
    case STATE_DISPENSING:
      // Updated in runDispensingLoop()
      break;
    case STATE_FAULT:
      lcd.print(F("!! FAULT ACTIVE "));
      break;
    case STATE_LOW_SUPPLY:
      lcd.print(F("LOW SUPPLY WARN "));
      break;
  }
}

// ────────────────────────────────────────────────────────────
void publishHeartbeat() {
  Serial.print(F("[MQTT] publish → hospital/supply/heartbeat : "));
  Serial.print(F("{node:SUPPLY, state:"));
  Serial.print(currentState);
  Serial.print(F(", water_pct:"));
  Serial.print(waterLevelPct, 1);
  Serial.print(F(", chem_ml:"));
  Serial.print(chemRemainingML, 0);
  Serial.println(F("}"));
}

// ────────────────────────────────────────────────────────────
void publishCycleLog(bool success, float volDispensed) {
  /*
   * Publishes cycle completion data to central controller.
   * Central controller logs this to database for reports.
   * Requirement 8: dispensing data collection
   */
  Serial.println(F("[MQTT] publish → hospital/supply/cycle_log :"));
  Serial.print(F("  success: "));      Serial.println(success);
  Serial.print(F("  volume_mL: "));    Serial.println(volDispensed, 1);
  Serial.print(F("  water_mL: "));     Serial.println(waterFlowTotal, 1);
  Serial.print(F("  chem_mL: "));      Serial.println(chemFlowTotal, 1);
  Serial.print(F("  final_ppm: "));    Serial.println(actualPPM, 0);
  Serial.print(F("  duration_ms: "));  Serial.println(millis() - cycleStartTime);
  Serial.print(F("  chem_remain_mL: ")); Serial.println(chemRemainingML, 0);
}

// ────────────────────────────────────────────────────────────
void checkSerialCommands() {
  /*
   * Simulates MQTT command reception via Serial Monitor.
   * In real system: PubSubClient MQTT callback.
   *
   * Commands:
   *   DISPENSE        → start dispensing cycle
   *   CLEAR_FAULT     → reset fault state
   *   STATUS          → print system status
   *   SET_WATER:XX    → set water level % (testing)
   *   SET_CHEM:XXXX   → set chemical level mL (testing)
   */

  if (Serial.available()) {
    String cmd = Serial.readStringUntil('\n');
    cmd.trim();
    cmd.toUpperCase();

    Serial.print(F("[CMD] Received: "));
    Serial.println(cmd);

    if (cmd == "DISPENSE") {
      startDispensingCycle(DISPENSE_TARGET_ML);

    } else if (cmd == "CLEAR_FAULT") {
      clearFault();

    } else if (cmd == "STATUS") {
      Serial.println(F("─── SYSTEM STATUS ───────────────"));
      Serial.print(F("  State: "));        Serial.println(currentState);
      Serial.print(F("  Water level: "));  Serial.print(waterLevelPct, 1); Serial.println(F("%"));
      Serial.print(F("  Chem remain: ")); Serial.print(chemRemainingML, 0); Serial.println(F(" mL"));
      Serial.print(F("  Fault active: ")); Serial.println(faultActive);
      Serial.println(F("─────────────────────────────────"));

    } else if (cmd.startsWith("SET_WATER:")) {
      waterLevelPct = cmd.substring(10).toFloat();
      Serial.print(F("[SIM] Water level set to: "));
      Serial.println(waterLevelPct);

    } else if (cmd.startsWith("SET_CHEM:")) {
      chemRemainingML = cmd.substring(9).toFloat();
      Serial.print(F("[SIM] Chemical level set to: "));
      Serial.println(chemRemainingML);

    } else {
      Serial.println(F("[CMD] Unknown command."));
      Serial.println(F("  Valid: DISPENSE, CLEAR_FAULT, STATUS, SET_WATER:XX, SET_CHEM:XXXX"));
    }
  }
}