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

// WiFi and MQTT Configuration
const char* ssid = "Wokwi-GUEST";
const char* password = "";
const char* mqtt_server = "broker.hivemq.com";
const char* mqtt_client_id = "WaterTankSystem_001";

// Pin Definitions - Updated to match diagram.json
// LEDs
#define LED_RED 2        // Connection status/Error LED
#define LED_GREEN_1 4    // Stage 1: System Check
#define LED_GREEN_2 5    // Stage 2: Water Analysis
#define LED_GREEN_3 18   // Stage 3: Distribution Control
#define LED_GREEN_4 16   // Stage 4: System Ready/Active

// I2C LCD Display (20x4) - Connected through PCF8574 I2C backpack
#define LCD_SDA 21       // I2C Data
#define LCD_SCL 22       // I2C Clock

// Sensors - Updated pin assignments from diagram
#define TRIG_PIN 19      // HC-SR04 Ultrasonic Trigger
#define ECHO_PIN 17      // HC-SR04 Ultrasonic Echo
#define FLOW_SENSOR_PIN 23  // Flow meter pulse input (simulated with button)
#define LEAK_SENSOR_PIN 15  // Leak detection (slide switch)
#define TDS_SENSOR_PIN 34   // TDS sensor (potentiometer)

// Outputs - Updated pin assignments from diagram
#define PUMP_RELAY_PIN 25   // Water pump control relay
#define VALVE_RELAY_PIN 26  // Solenoid valve control relay
#define BUZZER_PIN 27       // Alert buzzer

// Manual Controls - Updated pin assignments from diagram
#define BUTTON_START 32     // Start process button (green)
#define BUTTON_STOP 33      // Emergency stop button (red)

// System Constants
#define TANK_HEIGHT_CM 100     // Tank height in cm
#define MIN_WATER_LEVEL 20     // Minimum safe water level (cm from bottom)
#define MAX_WATER_LEVEL 90     // Maximum water level (cm from bottom)
#define FLOW_PULSE_PER_LITER 450  // Flow sensor calibration

// Global Variables
WiFiClient espClient;
PubSubClient client(espClient);
LiquidCrystal_I2C lcd(0x27, 20, 4); // I2C address 0x27, 20 columns, 4 rows

// System State
enum SystemState {
  SYSTEM_IDLE,
  SYSTEM_CHECK,
  WATER_ANALYSIS,
  DISTRIBUTION_CONTROL,
  SYSTEM_READY,
  ERROR_STATE,
  EMERGENCY_STOP
};

SystemState currentState = SYSTEM_IDLE;
bool systemEnabled = false;
bool emergencyStop = false;

// Sensor Data
struct SensorData {
  float waterLevel;      // cm from bottom
  float flowRate;        // L/min
  bool leakDetected;
  int tdsValue;          // ppm
  float temperature;     // °C
};

SensorData sensors;

// LCD Display tracking to prevent flickering
String lcdLine[4] = {"", "", "", ""};
bool lcdNeedsUpdate = true;

// Timing Variables
unsigned long lastSensorRead = 0;
unsigned long lastMqttPublish = 0;
unsigned long lastConnectionCheck = 0;
unsigned long lastButtonCheck = 0;
unsigned long lastLcdUpdate = 0;
volatile unsigned long flowPulseCount = 0;
unsigned long lastFlowTime = 0;

// Connection status
bool wifiConnected = false;
bool mqttConnected = false;

void setup() {
  Serial.begin(115200);
  delay(1000); // Give serial time to initialize
  Serial.println("Smart Water Tank System Starting...");
  
  // Initialize pins first
  initializePins();
  
  // Initialize LCD with error handling
  initializeLCD();
  
  // Initialize sensors
  initializeSensors();
  
  // Connect to WiFi
  setupWiFi();
  
  // Setup MQTT
  client.setServer(mqtt_server, 1883);
  client.setCallback(mqttCallback);
  
  // Initial system check
  performSystemCheck();
  
  // Force initial LCD update
  lcdNeedsUpdate = true;
  updateLcdDisplay();
  
  Serial.println("System initialization complete");
}

void loop() {
  unsigned long currentTime = millis();
  
  // Connection management
  if (currentTime - lastConnectionCheck > 5000) {
    checkConnections();
    lastConnectionCheck = currentTime;
  }
  
  // Read sensors every 1 second
  if (currentTime - lastSensorRead > 1000) {
    readSensors();
    lastSensorRead = currentTime;
  }
  
  // Check manual buttons
  if (currentTime - lastButtonCheck > 100) {
    checkButtons();
    lastButtonCheck = currentTime;
  }
  
  // Publish MQTT data every 5 seconds
  if (currentTime - lastMqttPublish > 5000) {
    publishSensorData();
    lastMqttPublish = currentTime;
  }
  
  // Update LCD every 2 seconds
  if (currentTime - lastLcdUpdate > 2000) {
    updateLcdDisplay();
    lastLcdUpdate = currentTime;
  }
  
  // State machine execution
  executeStateMachine();
  
  // Update LED status
  updateStatusLEDs();
  
  // Safety checks
  performSafetyChecks();
  
  client.loop();
  delay(50);
}

void initializeLCD() {
  Serial.println("Initializing LCD...");
  
  // Initialize I2C with specific pins
  Wire.begin(LCD_SDA, LCD_SCL);
  
  // Scan for I2C devices
  Serial.println("Scanning I2C devices...");
  bool deviceFound = false;
  for (byte address = 1; address < 127; address++) {
    Wire.beginTransmission(address);
    if (Wire.endTransmission() == 0) {
      Serial.print("I2C device found at address 0x");
      if (address < 16) Serial.print("0");
      Serial.println(address, HEX);
      deviceFound = true;
    }
  }
  
  if (!deviceFound) {
    Serial.println("No I2C devices found! Check wiring.");
  }
  
  // Try different common I2C addresses
  byte lcdAddresses[] = {0x27, 0x3F, 0x38, 0x20};
  bool lcdInitialized = false;
  
  for (int i = 0; i < 4 && !lcdInitialized; i++) {
    Serial.print("Trying LCD address: 0x");
    Serial.println(lcdAddresses[i], HEX);
    
    lcd = LiquidCrystal_I2C(lcdAddresses[i], 20, 4);
    lcd.init();
    
    // Test if LCD responds
    Wire.beginTransmission(lcdAddresses[i]);
    if (Wire.endTransmission() == 0) {
      lcd.backlight();
      lcd.clear();
      lcd.setCursor(0, 0);
      lcd.print("Water Tank System");
      lcd.setCursor(0, 1);
      lcd.print("LCD Initialized");
      
      // Create custom degree symbol (°)
      byte degreeChar[8] = {
        0b00110,
        0b01001,
        0b01001,
        0b00110,
        0b00000,
        0b00000,
        0b00000,
        0b00000
      };
      lcd.createChar(0, degreeChar);
      
      Serial.print("LCD initialized at address: 0x");
      Serial.println(lcdAddresses[i], HEX);
      lcdInitialized = true;
      delay(2000);
    }
  }
  
  if (!lcdInitialized) {
    Serial.println("LCD initialization failed!");
    // Continue without LCD
  }
}

void initializePins() {
  // LED pins - matching diagram connections
  pinMode(LED_RED, OUTPUT);        // D2
  pinMode(LED_GREEN_1, OUTPUT);    // D4
  pinMode(LED_GREEN_2, OUTPUT);    // D5
  pinMode(LED_GREEN_3, OUTPUT);    // D18
  pinMode(LED_GREEN_4, OUTPUT);    // D16
  
  // Sensor pins - matching diagram connections
  pinMode(TRIG_PIN, OUTPUT);       // D19 to ultrasonic TRIG
  pinMode(ECHO_PIN, INPUT);        // D17 to ultrasonic ECHO
  pinMode(LEAK_SENSOR_PIN, INPUT_PULLUP); // D15 to slide switch
  pinMode(TDS_SENSOR_PIN, INPUT);  // D34 to potentiometer SIG
  
  // Output pins - matching diagram connections
  pinMode(PUMP_RELAY_PIN, OUTPUT); // D25 to pump relay IN
  pinMode(VALVE_RELAY_PIN, OUTPUT); // D26 to valve relay IN
  pinMode(BUZZER_PIN, OUTPUT);     // D27 to buzzer
  
  // Button pins - matching diagram connections
  pinMode(BUTTON_START, INPUT_PULLUP); // D32 to green button
  pinMode(BUTTON_STOP, INPUT_PULLUP);  // D33 to red button
  
  // Flow sensor interrupt - matching diagram connections
  pinMode(FLOW_SENSOR_PIN, INPUT_PULLUP); // D23 to blue button (flow simulation)
  attachInterrupt(digitalPinToInterrupt(FLOW_SENSOR_PIN), flowPulseCounter, FALLING);
  
  // Initialize outputs to safe state
  digitalWrite(PUMP_RELAY_PIN, LOW);
  digitalWrite(VALVE_RELAY_PIN, LOW);
  digitalWrite(BUZZER_PIN, LOW);
  
  // Turn off all LEDs initially
  digitalWrite(LED_RED, LOW);
  digitalWrite(LED_GREEN_1, LOW);
  digitalWrite(LED_GREEN_2, LOW);
  digitalWrite(LED_GREEN_3, LOW);
  digitalWrite(LED_GREEN_4, LOW);
  
  Serial.println("Pins initialized");
}

void initializeSensors() {
  sensors.waterLevel = 50.0; // Default safe value
  sensors.flowRate = 0;
  sensors.leakDetected = false;
  sensors.tdsValue = 0;
  sensors.temperature = 25.0;
  
  Serial.println("Sensors initialized");
}

void setupWiFi() {
  Serial.print("Connecting to WiFi");
  
  // Display on LCD without flickering
  updateLcdLine(0, "Water Tank System");
  updateLcdLine(1, "Connecting WiFi...  ");
  updateLcdLine(2, "");
  updateLcdLine(3, "");
  
  WiFi.begin(ssid, password);
  
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    digitalWrite(LED_RED, HIGH);
    delay(250);
    digitalWrite(LED_RED, LOW);
    delay(250);
    Serial.print(".");
    attempts++;
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    wifiConnected = true;
    Serial.println();
    Serial.print("WiFi connected! IP: ");
    Serial.println(WiFi.localIP());
    
    updateLcdLine(1, "WiFi Connected      ");
    String ipLine = "IP: " + WiFi.localIP().toString();
    updateLcdLine(2, ipLine);
  } else {
    Serial.println("\nWiFi connection failed!");
    wifiConnected = false;
    updateLcdLine(1, "WiFi Failed!        ");
  }
  delay(2000);
}

void checkConnections() {
  // Check WiFi
  if (WiFi.status() != WL_CONNECTED) {
    wifiConnected = false;
    WiFi.reconnect();
  } else {
    wifiConnected = true;
  }
  
  // Check MQTT
  if (!client.connected()) {
    mqttConnected = false;
    reconnectMQTT();
  } else {
    mqttConnected = true;
  }
}

void reconnectMQTT() {
  if (!wifiConnected) return;
  
  int attempts = 0;
  while (!client.connected() && attempts < 3) {
    Serial.print("Attempting MQTT connection...");
    
    if (client.connect(mqtt_client_id)) {
      Serial.println("connected");
      mqttConnected = true;
      
      // Subscribe to control topics
      client.subscribe("watertank/control/start");
      client.subscribe("watertank/control/stop");
      client.subscribe("watertank/control/emergency");
      client.subscribe("watertank/config/set_level");
      
      // Publish connection status
      client.publish("watertank/status/online", "true", true);
      
    } else {
      Serial.print("failed, rc=");
      Serial.print(client.state());
      Serial.println(" try again in 5 seconds");
      attempts++;
    }
  }
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  String message;
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  
  Serial.print("MQTT received [");
  Serial.print(topic);
  Serial.print("]: ");
  Serial.println(message);
  
  // Handle commands
  if (String(topic) == "watertank/control/start") {
    if (!emergencyStop) {
      systemEnabled = true;
      currentState = SYSTEM_CHECK;
      Serial.println("System start command received");
    }
  }
  else if (String(topic) == "watertank/control/stop") {
    systemEnabled = false;
    currentState = SYSTEM_IDLE;
    stopAllOutputs();
    Serial.println("System stop command received");
  }
  else if (String(topic) == "watertank/control/emergency") {
    emergencyStop = true;
    currentState = EMERGENCY_STOP;
    stopAllOutputs();
    Serial.println("Emergency stop activated!");
  }
}

void readSensors() {
  // Read water level (HC-SR04) with improved timeout handling
  sensors.waterLevel = readWaterLevel();
  
  // Read flow rate
  sensors.flowRate = calculateFlowRate();
  
  // Read leak sensor - slide switch (LOW when switch is moved to simulate leak)
  sensors.leakDetected = !digitalRead(LEAK_SENSOR_PIN);
  
  // Read TDS sensor - potentiometer value converted to ppm
  int rawTDS = analogRead(TDS_SENSOR_PIN);
  sensors.tdsValue = map(rawTDS, 0, 4095, 0, 1000); // Convert to ppm
  
  // Simulate temperature (in real system, use DS18B20)
  sensors.temperature = 25.0 + (random(-50, 50) / 10.0);
  
  // Debug output
  Serial.print("Water Level: ");
  Serial.print(sensors.waterLevel);
  Serial.print(" cm, TDS: ");
  Serial.print(sensors.tdsValue);
  Serial.print(" ppm, Flow: ");
  Serial.print(sensors.flowRate);
  Serial.println(" L/min");
}

float readWaterLevel() {
  // Multiple attempt strategy for ultrasonic sensor
  float validReadings[3];
  int validCount = 0;
  
  for (int attempt = 0; attempt < 3 && validCount < 2; attempt++) {
    // Clear trigger
    digitalWrite(TRIG_PIN, LOW);
    delayMicroseconds(5);
    
    // Send pulse
    digitalWrite(TRIG_PIN, HIGH);
    delayMicroseconds(10);
    digitalWrite(TRIG_PIN, LOW);
    
    // Read echo with longer timeout for Wokwi simulation
    long duration = pulseIn(ECHO_PIN, HIGH, 50000); // Increased timeout to 50ms
    
    if (duration > 0) {
      float distance = (duration * 0.034) / 2.0; // Convert to cm
      float waterLevel = TANK_HEIGHT_CM - distance;
      
      // Validate reading
      if (waterLevel >= 0 && waterLevel <= TANK_HEIGHT_CM) {
        validReadings[validCount] = waterLevel;
        validCount++;
        Serial.print("Valid ultrasonic reading ");
        Serial.print(validCount);
        Serial.print(": ");
        Serial.print(waterLevel);
        Serial.println(" cm");
      } else {
        Serial.print("Invalid ultrasonic reading: ");
        Serial.print(waterLevel);
        Serial.println(" cm (out of range)");
      }
    } else {
      Serial.print("Ultrasonic timeout on attempt ");
      Serial.println(attempt + 1);
    }
    
    delay(100); // Small delay between attempts
  }
  
  // Process valid readings
  if (validCount > 0) {
    // Use average if multiple readings
    float sum = 0;
    for (int i = 0; i < validCount; i++) {
      sum += validReadings[i];
    }
    float newLevel = sum / validCount;
    
    // Apply simple filtering to prevent wild swings
    if (abs(newLevel - sensors.waterLevel) > 10) {
      // Large change - use weighted average
      return (sensors.waterLevel * 0.7) + (newLevel * 0.3);
    } else {
      return newLevel;
    }
  } else {
    // No valid readings - use simulated value for Wokwi
    Serial.println("Using simulated water level for Wokwi");
    // Simulate water level based on pump status
    static float simulatedLevel = 50.0;
    if (digitalRead(PUMP_RELAY_PIN) == HIGH) {
      simulatedLevel += 0.5; // Filling
      if (simulatedLevel > MAX_WATER_LEVEL) {
        simulatedLevel = MAX_WATER_LEVEL;
      }
    } else {
      simulatedLevel -= 0.1; // Natural consumption
      if (simulatedLevel < MIN_WATER_LEVEL) {
        simulatedLevel = MIN_WATER_LEVEL;
      }
    }
    return simulatedLevel;
  }
}

float calculateFlowRate() {
  unsigned long currentTime = millis();
  float timeInterval = (currentTime - lastFlowTime) / 1000.0; // seconds
  
  if (timeInterval >= 1.0) {
    float pulseRate = flowPulseCount / timeInterval;
    float flowRate = pulseRate / FLOW_PULSE_PER_LITER * 60; // L/min
    
    flowPulseCount = 0;
    lastFlowTime = currentTime;
    
    return flowRate;
  }
  
  return sensors.flowRate; // Return previous reading
}

IRAM_ATTR void flowPulseCounter() {
  flowPulseCount++;
}

void checkButtons() {
  static bool lastStartState = HIGH;
  static bool lastStopState = HIGH;
  
  bool currentStartState = digitalRead(BUTTON_START);
  bool currentStopState = digitalRead(BUTTON_STOP);
  
  // Start button pressed (active LOW due to pull-up)
  if (lastStartState == HIGH && currentStartState == LOW) {
    if (!emergencyStop) {
      systemEnabled = true;
      currentState = SYSTEM_CHECK;
      Serial.println("Manual start button pressed");
    }
  }
  
  // Stop button pressed (active LOW due to pull-up)
  if (lastStopState == HIGH && currentStopState == LOW) {
    emergencyStop = true;
    currentState = EMERGENCY_STOP;
    stopAllOutputs();
    Serial.println("Manual emergency stop pressed!");
  }
  
  lastStartState = currentStartState;
  lastStopState = currentStopState;
}

void executeStateMachine() {
  static unsigned long stateStartTime = 0;
  static SystemState lastState = SYSTEM_IDLE;
  
  if (currentState != lastState) {
    stateStartTime = millis();
    lastState = currentState;
    Serial.print("State changed to: ");
    Serial.println(stateToString(currentState));
  }
  
  unsigned long stateTime = millis() - stateStartTime;
  
  switch (currentState) {
    case SYSTEM_IDLE:
      handleIdleState();
      break;
      
    case SYSTEM_CHECK:
      if (handleSystemCheckState(stateTime)) {
        currentState = WATER_ANALYSIS;
      }
      break;
      
    case WATER_ANALYSIS:
      if (handleWaterAnalysisState(stateTime)) {
        currentState = DISTRIBUTION_CONTROL;
      }
      break;
      
    case DISTRIBUTION_CONTROL:
      handleDistributionControlState();
      break;
      
    case SYSTEM_READY:
      handleSystemReadyState();
      break;
      
    case ERROR_STATE:
      handleErrorState();
      break;
      
    case EMERGENCY_STOP:
      handleEmergencyStopState();
      break;
  }
}

void handleIdleState() {
  stopAllOutputs();
  
  if (systemEnabled && !emergencyStop) {
    currentState = SYSTEM_CHECK;
  }
}

bool handleSystemCheckState(unsigned long stateTime) {
  // Stage 1: System Check
  digitalWrite(LED_GREEN_1, HIGH);
  
  // Check all sensors
  bool sensorsOK = true;
  
  // Check water level sensor
  if (sensors.waterLevel < 0 || sensors.waterLevel > TANK_HEIGHT_CM) {
    sensorsOK = false;
    Serial.println("Water level sensor error");
  }
  
  // Check for critical water levels
  if (sensors.waterLevel < MIN_WATER_LEVEL) {
    Serial.println("Warning: Water level low but continuing check");
    // Don't fail system check for low water - this is normal
  }
  
  // Minimum check time of 3 seconds
  if (stateTime > 3000 && sensorsOK) {
    digitalWrite(LED_GREEN_1, LOW);
    return true; // Move to next stage
  } else if (stateTime > 10000) {
    // Timeout or error
    currentState = ERROR_STATE;
    return false;
  }
  
  return false;
}

bool handleWaterAnalysisState(unsigned long stateTime) {
  // Stage 2: Water Analysis
  digitalWrite(LED_GREEN_2, HIGH);
  
  // Analyze water quality
  bool waterQualityOK = true;
  
  if (sensors.tdsValue > 800) {
    waterQualityOK = false;
    Serial.println("Water TDS too high");
  }
  
  // Minimum analysis time of 2 seconds
  if (stateTime > 2000 && waterQualityOK) {
    digitalWrite(LED_GREEN_2, LOW);
    return true; // Move to next stage
  } else if (stateTime > 8000) {
    // Timeout or error
    currentState = ERROR_STATE;
    return false;
  }
  
  return false;
}

void handleDistributionControlState() {
  // Stage 3: Distribution Control
  digitalWrite(LED_GREEN_3, HIGH);
  
  // Control pump based on water level
  if (sensors.waterLevel < MAX_WATER_LEVEL && !sensors.leakDetected) {
    digitalWrite(PUMP_RELAY_PIN, HIGH);
    digitalWrite(VALVE_RELAY_PIN, HIGH);
  } else {
    digitalWrite(PUMP_RELAY_PIN, LOW);
    
    if (sensors.waterLevel >= MAX_WATER_LEVEL) {
      digitalWrite(VALVE_RELAY_PIN, LOW);
      digitalWrite(LED_GREEN_3, LOW);
      currentState = SYSTEM_READY;
      Serial.println("Tank full - moving to ready state");
    }
  }
  
  // Check if system should continue
  if (!systemEnabled || emergencyStop) {
    digitalWrite(LED_GREEN_3, LOW);
    currentState = SYSTEM_IDLE;
  }
}

void handleSystemReadyState() {
  // Stage 4: System Ready/Monitoring
  digitalWrite(LED_GREEN_4, HIGH);
  
  // Monitor water level and maintain
  if (sensors.waterLevel < (MAX_WATER_LEVEL - 5) && !sensors.leakDetected) {
    // Water level dropped, restart distribution
    currentState = DISTRIBUTION_CONTROL;
    digitalWrite(LED_GREEN_4, LOW);
  }
  
  // Check if system should stop
  if (!systemEnabled || emergencyStop) {
    digitalWrite(LED_GREEN_4, LOW);
    currentState = SYSTEM_IDLE;
  }
}

void handleErrorState() {
  stopAllOutputs();
  
  // Blink red LED for error
  static unsigned long lastBlink = 0;
  if (millis() - lastBlink > 500) {
    digitalWrite(LED_RED, !digitalRead(LED_RED));
    lastBlink = millis();
  }
  
  // Sound buzzer intermittently
  static unsigned long lastBuzz = 0;
  if (millis() - lastBuzz > 2000) {
    tone(BUZZER_PIN, 1000, 500);
    lastBuzz = millis();
  }
  
  // Reset error state manually or via MQTT
  if (!systemEnabled) {
    currentState = SYSTEM_IDLE;
  }
}

void handleEmergencyStopState() {
  stopAllOutputs();
  
  // Solid red LED
  digitalWrite(LED_RED, HIGH);
  
  // Continuous buzzer for 3 seconds
  static bool buzzerSounded = false;
  static unsigned long buzzerStartTime = 0;
  
  if (!buzzerSounded) {
    tone(BUZZER_PIN, 2000);
    buzzerSounded = true;
    buzzerStartTime = millis();
  }
  
  // Stop buzzer after 3 seconds
  if (millis() - buzzerStartTime > 3000) {
    noTone(BUZZER_PIN);
  }
  
  // Manual reset required - check for reset button combination
  if (digitalRead(BUTTON_START) == LOW && digitalRead(BUTTON_STOP) == LOW) {
    // Both buttons pressed - reset emergency stop
    emergencyStop = false;
    systemEnabled = false;
    currentState = SYSTEM_IDLE;
    buzzerSounded = false;
    digitalWrite(LED_RED, LOW);
    Serial.println("Emergency stop reset");
  }
}

void performSafetyChecks() {
  // Leak detection
  if (sensors.leakDetected) {
    Serial.println("LEAK DETECTED!");
    stopAllOutputs();
    currentState = ERROR_STATE;
    if (mqttConnected) {
      client.publish("watertank/alert/leak", "CRITICAL");
    }
  }
  
  // Low water level protection
  if (sensors.waterLevel < MIN_WATER_LEVEL && digitalRead(PUMP_RELAY_PIN)) {
    Serial.println("Low water level - stopping pump");
    digitalWrite(PUMP_RELAY_PIN, LOW);
    if (mqttConnected) {
      client.publish("watertank/alert/low_level", "WARNING");
    }
  }
  
  // Overflow protection
  if (sensors.waterLevel > MAX_WATER_LEVEL) {
    Serial.println("High water level - closing valve");
    digitalWrite(VALVE_RELAY_PIN, LOW);
    if (mqttConnected) {
      client.publish("watertank/alert/high_level", "WARNING");
    }
  }
}

void updateStatusLEDs() {
  static unsigned long lastUpdate = 0;
  static bool ledState = false;
  
  if (millis() - lastUpdate > 1000) {
    lastUpdate = millis();
    ledState = !ledState;
    
    // Connection status LED
    if (!wifiConnected || !mqttConnected) {
      digitalWrite(LED_RED, ledState); // Blink for connection issues
    } else if (currentState != ERROR_STATE && currentState != EMERGENCY_STOP) {
      digitalWrite(LED_RED, LOW); // Solid off when connected
    }
  }
}

void publishSensorData() {
  if (!mqttConnected) return;
  
  // Create JSON payload
  StaticJsonDocument<512> doc;
  doc["water_level"] = sensors.waterLevel;
  doc["flow_rate"] = sensors.flowRate;
  doc["tds"] = sensors.tdsValue;
  doc["temperature"] = sensors.temperature;
  doc["leak_detected"] = sensors.leakDetected;
  doc["pump_status"] = digitalRead(PUMP_RELAY_PIN);
  doc["valve_status"] = digitalRead(VALVE_RELAY_PIN);
  doc["system_state"] = stateToString(currentState);
  doc["timestamp"] = millis();
  
  char buffer[512];
  serializeJson(doc, buffer);
  
  client.publish("watertank/telemetry", buffer);
  
  // Individual topic updates
  client.publish("watertank/status/level", String(sensors.waterLevel).c_str());
  client.publish("watertank/status/flow", String(sensors.flowRate).c_str());
  client.publish("watertank/status/tds", String(sensors.tdsValue).c_str());
}

void stopAllOutputs() {
  digitalWrite(PUMP_RELAY_PIN, LOW);
  digitalWrite(VALVE_RELAY_PIN, LOW);
  noTone(BUZZER_PIN);
  digitalWrite(LED_GREEN_1, LOW);
  digitalWrite(LED_GREEN_2, LOW);
  digitalWrite(LED_GREEN_3, LOW);
  digitalWrite(LED_GREEN_4, LOW);
}

void performSystemCheck() {
  Serial.println("Performing initial system check...");
  
  updateLcdLine(0, "System Check...");
  updateLcdLine(1, "");
  updateLcdLine(2, "");
  updateLcdLine(3, "");
  
  // Test all LEDs
  updateLcdLine(1, "Testing LEDs...");
  digitalWrite(LED_RED, HIGH);
  digitalWrite(LED_GREEN_1, HIGH);
  digitalWrite(LED_GREEN_2, HIGH);
  digitalWrite(LED_GREEN_3, HIGH);
  digitalWrite(LED_GREEN_4, HIGH);
  delay(1000);
  
  digitalWrite(LED_RED, LOW);
  digitalWrite(LED_GREEN_1, LOW);
  digitalWrite(LED_GREEN_2, LOW);
  digitalWrite(LED_GREEN_3, LOW);
  digitalWrite(LED_GREEN_4, LOW);
  
  // Test buzzer
  updateLcdLine(2, "Testing Buzzer...");
  tone(BUZZER_PIN, 1000, 500);
  delay(600);
  
  updateLcdLine(3, "Check Complete");
  delay(1000);
  
  Serial.println("Initial system check complete");
}

// Function to update LCD line only if content has changed (prevents flickering)
void updateLcdLine(int line, String content) {
  // Pad or trim content to exactly 20 characters to prevent artifacts
  while (content.length() < 20) {
    content += " ";
  }
  if (content.length() > 20) {
    content = content.substring(0, 20);
  }
  
  // Only update if content has changed
  if (lcdLine[line] != content) {
    lcdLine[line] = content;
    lcd.setCursor(0, line);
    lcd.print(content);
  }
}

void updateLcdDisplay() {
  // Line 1: System status
  String line1 = "Status: " + String(stateToString(currentState));
  updateLcdLine(0, line1);
  
  // Line 2: Water level and TDS with proper spacing
  String line2 = "Level:" + String(sensors.waterLevel, 1) + "cm TDS:" + String(sensors.tdsValue);
  updateLcdLine(1, line2);
  
  // Line 3: Flow rate and temperature with degree symbol (custom character 0)
  String line3 = "Flow:" + String(sensors.flowRate, 1) + "L/m Temp:" + String(sensors.temperature, 1);
  line3 += char(0); // Custom degree symbol
  line3 += "C";
  updateLcdLine(2, line3);
  
  // Line 4: Connection status and alerts
  String line4;
  if (sensors.leakDetected) {
    line4 = "*** LEAK ALERT ***";
  } else if (!wifiConnected) {
    line4 = "WiFi: Disconnected";
  } else if (!mqttConnected) {
    line4 = "MQTT: Disconnected";
  } else {
    line4 = "Pump:" + String(digitalRead(PUMP_RELAY_PIN) ? "ON " : "OFF");
    line4 += " Valve:" + String(digitalRead(VALVE_RELAY_PIN) ? "ON " : "OFF");
  }
  updateLcdLine(3, line4);
}

const char* stateToString(SystemState state) {
  switch (state) {
    case SYSTEM_IDLE:
      return "IDLE";
    case SYSTEM_CHECK:
      return "CHECK";
    case WATER_ANALYSIS:
      return "ANALYSIS";
    case DISTRIBUTION_CONTROL:
      return "CONTROL";
    case SYSTEM_READY:
      return "READY";
    case ERROR_STATE:
      return "ERROR";
    case EMERGENCY_STOP:
      return "E-STOP";
    default:
      return "UNKNOWN";
  }
}