#include <WiFi.h>
#include <PubSubClient.h>
#include <ArduinoJson.h>
#include <LiquidCrystal_I2C.h>
#include <DHT.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 = "IndustrialIoT_001";

// Pin Definitions - Updated to match Industrial IoT diagram
// Sensors
#define DHT_PIN 4            // DHT22 Environment sensor
#define DHT_TYPE DHT22
#define PRESSURE_PIN A0      // Pressure sensor (potentiometer 1)
#define FLOW_PIN A3          // Flow rate sensor (potentiometer 2)  
#define VIBRATION_PIN A6     // Vibration analyzer (potentiometer 3)
#define LIGHT_PIN A7         // Light sensor (LDR)
#define MOTION_PIN 13        // PIR Motion detector

// I2C LCD Display - Fixed address for production
#define LCD_SDA 21           // I2C Data
#define LCD_SCL 22           // I2C Clock
#define LCD_I2C_ADDR 0x27    // Fixed I2C address for production

// Control Relays
#define MOTOR_RELAY_PIN 25   // Motor Controller
#define PUMP_RELAY_PIN 26    // Pump System  
#define VALVE_RELAY_PIN 27   // Valve Control
#define COOLING_RELAY_PIN 14 // Cooling System

// Status LEDs
#define LED_SYSTEM_ONLINE 2  // Green - System Online
#define LED_ALERT 16         // Red - Alert Status
#define LED_PROCESS_ACTIVE 17 // Blue - Process Active
#define LED_POWER_STATUS 18  // Yellow - Power Status

// Manual Controls
#define BUTTON_EMERGENCY 15  // Emergency Stop (red button)
#define BUTTON_START 32      // System Start (green button)

// Buzzer
#define BUZZER_PIN 19        // Alarm buzzer

// System Constants
#define PRESSURE_MIN 0       // Minimum pressure (bar)
#define PRESSURE_MAX 10      // Maximum pressure (bar)
#define FLOW_MIN 0           // Minimum flow rate (L/min)
#define FLOW_MAX 100         // Maximum flow rate (L/min)
#define VIBRATION_THRESHOLD 80 // Vibration alert threshold
#define TEMP_HIGH_THRESHOLD 35.0 // High temperature threshold
#define EMERGENCY_RESET_TIME 3000 // Time to hold start button to reset emergency

// Global Variables
WiFiClient espClient;
PubSubClient client(espClient);
LiquidCrystal_I2C lcd(LCD_I2C_ADDR, 16, 2); // Fixed I2C address

// Sensors
DHT dht(DHT_PIN, DHT_TYPE);

// System State - Simplified state machine
enum SystemState {
  SYSTEM_IDLE,
  SYSTEM_STARTUP,
  SYSTEM_ACTIVE,    // Combined running and monitoring
  SYSTEM_ALERT,
  EMERGENCY_STOP
};

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

// MQTT Control Flags for sensors
bool motionSensorEnabled = true;
bool lightSensorEnabled = true;

// LED Manual Override flags
bool ledRedManualOverride = false;
bool ledBlueManualOverride = false;
bool ledRedManualState = false;
bool ledBlueManualState = false;

// Sensor Data Structure
struct SensorData {
  float environmentTemp;     // DHT22 temperature
  float environmentHumidity; // DHT22 humidity  
  float pressure;            // System pressure (bar)
  float flowRate;            // Flow rate (L/min)
  float vibration;           // Vibration level (%)
  int lightLevel;            // Light intensity (0-100%)
  bool motionDetected;       // Motion detection
};

SensorData sensors;

// Control States
struct ControlStates {
  bool motorRunning;
  bool pumpActive;
  bool valveOpen;
  bool coolingActive;
};

ControlStates controls;

// LCD Display tracking to prevent flickering
String lcdLine[2] = {"", ""};

// Timing Variables - Non-blocking timers (SLOWER UPDATES)
unsigned long lastSensorRead = 0;
unsigned long lastMqttPublish = 0;
unsigned long lastConnectionCheck = 0;
unsigned long lastButtonCheck = 0;
unsigned long lastLcdUpdate = 0;
unsigned long lastStateCheck = 0;
unsigned long lastSerialOutput = 0;

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

// Alert tracking
bool temperatureAlert = false;
bool pressureAlert = false;
bool vibrationAlert = false;
unsigned long lastAlertTime = 0;

// Emergency stop reset timing
unsigned long emergencyResetStartTime = 0;
bool emergencyResetInProgress = false;

void setup() {
  Serial.begin(115200);
  delay(1000);
  
  // Pretty header for serial monitor
  Serial.println("\n" + String('=', 50));
  Serial.println("🏭 INDUSTRIAL IoT CONTROL SYSTEM 🏭");
  Serial.println(String('=', 50));
  Serial.println("Version: 2.2 | Status: INITIALIZING...");
  Serial.println(String('-', 50));
  
  // Initialize pins
  initializePins();
  
  // Initialize sensors
  initializeSensors();
  
  // Initialize LCD
  initializeLCD();
  
  // Connect to WiFi
  setupWiFi();
  
  // Setup MQTT
  client.setServer(mqtt_server, 1883);
  client.setCallback(mqttCallback);
  
  // Initial system check
  performSystemCheck();
  
  // Turn on Blue LED automatically at startup
  digitalWrite(LED_PROCESS_ACTIVE, HIGH);
  processActive = true;
  
  Serial.println(String('-', 50));
  Serial.println("✅ SYSTEM READY - Blue LED Auto-Activated");
  Serial.println(String('=', 50) + "\n");
}

void loop() {
  unsigned long currentTime = millis();
  
  // Connection management every 15 seconds (slower)
  if (currentTime - lastConnectionCheck > 15000) {
    checkConnections();
    lastConnectionCheck = currentTime;
  }
  
  // Read sensors every 6 seconds (slower)
  if (currentTime - lastSensorRead > 6000) {
    readSensors();
    lastSensorRead = currentTime;
  }
  
  // Check manual buttons every 100ms
  if (currentTime - lastButtonCheck > 100) {
    checkButtons();
    lastButtonCheck = currentTime;
  }
  
  // Publish MQTT data every 20 seconds (slower)
  if (currentTime - lastMqttPublish > 20000) {
    publishSensorData();
    lastMqttPublish = currentTime;
  }
  
  // Update LCD every 8 seconds (slower)
  if (currentTime - lastLcdUpdate > 8000) {
    updateLcdDisplay();
    lastLcdUpdate = currentTime;
  }
  
  // State machine execution every 3 seconds (slower)
  if (currentTime - lastStateCheck > 3000) {
    executeStateMachine();
    lastStateCheck = currentTime;
  }
  
  // Serial output every 12 seconds (slower serial updates)
  if (currentTime - lastSerialOutput > 12000) {
    printSystemStatus();
    lastSerialOutput = currentTime;
  }
  
  // Update status LEDs
  updateStatusLEDs();
  
  // Safety and alert checks
  performAlertChecks();
  
  client.loop();
  delay(150); // Increased delay
}

void printSystemStatus() {
  Serial.println("\n📊 SYSTEM STATUS");
  Serial.println(String('-', 30));
  
  // Compact system info
  Serial.print("🔧 State: ");
  Serial.print(stateToString(currentState));
  Serial.print(" | Enabled: ");
  Serial.print(systemEnabled ? "YES" : "NO");
  Serial.print(" | E-Stop: ");
  Serial.println(emergencyStop ? "ACTIVE" : "NORMAL");
  
  // Compact sensor readings
  Serial.print("🌡️ Env: ");
  Serial.print(sensors.environmentTemp, 1);
  Serial.print("°C ");
  Serial.print(sensors.environmentHumidity, 0);
  Serial.print("% | Press: ");
  Serial.print(sensors.pressure, 1);
  Serial.print("br | Flow: ");
  Serial.print(sensors.flowRate, 0);
  Serial.print("L/m | Vib: ");
  Serial.print(sensors.vibration, 0);
  Serial.print("% | Light: ");
  Serial.print(sensors.lightLevel);
  Serial.print("% | Motion: ");
  Serial.println(sensors.motionDetected ? "Y" : "N");
  
  // Compact control status
  Serial.print("⚙️ Controls: M:");
  Serial.print(controls.motorRunning ? "ON" : "OFF");
  Serial.print(" P:");
  Serial.print(controls.pumpActive ? "ON" : "OFF");
  Serial.print(" V:");
  Serial.print(controls.valveOpen ? "OPEN" : "SHUT");
  Serial.print(" C:");
  Serial.println(controls.coolingActive ? "ON" : "OFF");
                
  // LED Status with current states
  Serial.print("🚦 LED STATUS: Red LED: ");
  Serial.print(ledRedManualOverride ? (ledRedManualState ? "ON" : "OFF") : 
               (digitalRead(LED_ALERT) ? "ON" : "OFF"));
  Serial.print(" | Blue LED: ");
  Serial.println(ledBlueManualOverride ? (ledBlueManualState ? "ON" : "OFF") : 
                 (digitalRead(LED_PROCESS_ACTIVE) ? "ON" : "OFF"));
  
  // Compact connectivity
  Serial.print("🌐 WiFi: ");
  Serial.print(wifiConnected ? "OK" : "FAIL");
  Serial.print(" | MQTT: ");
  Serial.println(mqttConnected ? "OK" : "FAIL");
  
  // Active alerts (if any)
  if (temperatureAlert || pressureAlert || vibrationAlert) {
    Serial.print("⚠️ ALERTS: ");
    if (temperatureAlert) Serial.print("TEMP ");
    if (pressureAlert) Serial.print("PRESS ");
    if (vibrationAlert) Serial.print("VIB ");
    Serial.println();
  }
  
  Serial.println(String('-', 30));
}

void initializePins() {
  // Status LED pins
  pinMode(LED_SYSTEM_ONLINE, OUTPUT);
  pinMode(LED_ALERT, OUTPUT);
  pinMode(LED_PROCESS_ACTIVE, OUTPUT);
  pinMode(LED_POWER_STATUS, OUTPUT);
  
  // Control relay pins
  pinMode(MOTOR_RELAY_PIN, OUTPUT);
  pinMode(PUMP_RELAY_PIN, OUTPUT);
  pinMode(VALVE_RELAY_PIN, OUTPUT);
  pinMode(COOLING_RELAY_PIN, OUTPUT);
  
  // Button pins with internal pull-ups
  pinMode(BUTTON_EMERGENCY, INPUT_PULLUP);
  pinMode(BUTTON_START, INPUT_PULLUP);
  
  // Motion sensor pin
  pinMode(MOTION_PIN, INPUT);
  
  // Buzzer pin
  pinMode(BUZZER_PIN, OUTPUT);
  
  // Initialize outputs to safe state
  digitalWrite(MOTOR_RELAY_PIN, LOW);
  digitalWrite(PUMP_RELAY_PIN, LOW);
  digitalWrite(VALVE_RELAY_PIN, LOW);
  digitalWrite(COOLING_RELAY_PIN, LOW);
  digitalWrite(BUZZER_PIN, LOW);
  
  // Turn off all LEDs initially except power status
  digitalWrite(LED_SYSTEM_ONLINE, LOW);
  digitalWrite(LED_ALERT, LOW);
  digitalWrite(LED_PROCESS_ACTIVE, LOW);
  digitalWrite(LED_POWER_STATUS, HIGH); // Power on indicator
  
  Serial.println("📍 Hardware pins initialized successfully");
}

void initializeSensors() {
  // Initialize DHT22
  dht.begin();
  
  // Initialize sensor data with default values
  sensors.environmentTemp = 25.0;
  sensors.environmentHumidity = 50.0;
  sensors.pressure = 0.0;
  sensors.flowRate = 0.0;
  sensors.vibration = 0.0;
  sensors.lightLevel = 50;
  sensors.motionDetected = false;
  
  // Initialize control states
  controls.motorRunning = false;
  controls.pumpActive = false;
  controls.valveOpen = false;
  controls.coolingActive = false;
  
  Serial.println("🔬 Sensor array initialized successfully");
}

void initializeLCD() {
  Serial.println("🖥️ Initializing LCD display...");
  
  // Initialize I2C
  Wire.begin(LCD_SDA, LCD_SCL);
  
  // Use fixed address for production
  lcd.init();
  
  Wire.beginTransmission(LCD_I2C_ADDR);
  if (Wire.endTransmission() == 0) {
    lcd.backlight();
    lcd.clear();
    lcd.setCursor(0, 0);
    lcd.print("Industrial IoT");
    lcd.setCursor(0, 1);
    lcd.print("System Starting");
    
    Serial.print("📺 LCD initialized at address: 0x");
    Serial.println(LCD_I2C_ADDR, HEX);
    delay(2000);
  } else {
    Serial.println("❌ LCD not found at fixed address!");
  }
}

void setupWiFi() {
  Serial.println("📡 Connecting to WiFi network...");
  
  updateLcdLine(0, "Connecting WiFi");
  updateLcdLine(1, "Please wait...");
  
  WiFi.begin(ssid, password);
  
  int attempts = 0;
  while (WiFi.status() != WL_CONNECTED && attempts < 20) {
    digitalWrite(LED_ALERT, HIGH);
    delay(250);
    digitalWrite(LED_ALERT, LOW);
    delay(250);
    Serial.print(".");
    attempts++;
  }
  
  if (WiFi.status() == WL_CONNECTED) {
    wifiConnected = true;
    Serial.println();
    Serial.print("✅ WiFi connected successfully! IP: ");
    Serial.println(WiFi.localIP().toString());
    
    updateLcdLine(0, "WiFi Connected");
    updateLcdLine(1, WiFi.localIP().toString());
    digitalWrite(LED_SYSTEM_ONLINE, HIGH);
  } else {
    Serial.println("\n❌ WiFi connection failed!");
    wifiConnected = false;
    updateLcdLine(0, "WiFi Failed!");
    updateLcdLine(1, "Check connection");
  }
  delay(2000);
}

void checkConnections() {
  // Check WiFi
  if (WiFi.status() != WL_CONNECTED) {
    wifiConnected = false;
    digitalWrite(LED_SYSTEM_ONLINE, LOW);
    WiFi.reconnect();
  } else {
    wifiConnected = true;
    if (!emergencyStop) {
      digitalWrite(LED_SYSTEM_ONLINE, HIGH);
    }
  }
  
  // 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.println("🔄 Attempting MQTT connection...");
    
    if (client.connect(mqtt_client_id)) {
      Serial.println("✅ MQTT connected successfully");
      mqttConnected = true;
      
      // Subscribe to control topics
      client.subscribe("industrial/control/start");
      client.subscribe("industrial/control/stop");
      client.subscribe("industrial/control/emergency");
      client.subscribe("industrial/control/emergency_reset"); // NEW: Emergency reset command
      client.subscribe("industrial/control/motor");
      client.subscribe("industrial/control/pump");
      client.subscribe("industrial/control/valve");
      client.subscribe("industrial/control/cooling");
      
      // Subscribe to LED control topics
      client.subscribe("industrial/led/red");
      client.subscribe("industrial/led/blue");
      client.subscribe("industrial/led/green");
      client.subscribe("industrial/led/yellow");
      
      // Subscribe to sensor configuration topics
      client.subscribe("industrial/config/motion_enabled");
      client.subscribe("industrial/config/light_enabled");
      client.subscribe("industrial/config/+");
      
      // Publish connection status
      client.publish("industrial/status/online", "true", true);
      
      // Publish current configuration states
      publishConfigurationStates();
      
    } else {
      Serial.print("❌ MQTT connection failed, rc=");
      Serial.print(client.state());
      Serial.println(". Retrying in 5s...");
      attempts++;
    }
  }
}

void mqttCallback(char* topic, byte* payload, unsigned int length) {
  String message;
  for (int i = 0; i < length; i++) {
    message += (char)payload[i];
  }
  
  String topicStr = String(topic);
  
  // 📨 ALWAYS LOG INCOMING MQTT COMMANDS
  Serial.println("\n📨 MQTT COMMAND RECEIVED");
  Serial.print("   Topic: ");
  Serial.println(topicStr);
  Serial.print("   Message: ");
  Serial.println(message);
  Serial.println(String('-', 35));
  
  // System control commands
  if (topicStr == "industrial/control/start") {
    if (!emergencyStop) {
      systemEnabled = true;
      currentState = SYSTEM_STARTUP;
      Serial.println("✅ System START command executed");
    } else {
      Serial.println("❌ System START blocked - emergency stop active");
    }
  }
  else if (topicStr == "industrial/control/stop") {
    systemEnabled = false;
    processActive = false;
    currentState = SYSTEM_IDLE;
    stopAllOutputs();
    Serial.println("✅ System STOP command executed");
  }
  else if (topicStr == "industrial/control/emergency") {
    emergencyStop = true;
    currentState = EMERGENCY_STOP;
    stopAllOutputs();
    Serial.println("🚨 EMERGENCY STOP activated via MQTT!");
  }
  // NEW: Emergency reset command
  else if (topicStr == "industrial/control/emergency_reset") {
    if (emergencyStop) {
      emergencyStop = false;
      systemEnabled = false;
      currentState = SYSTEM_IDLE;
      emergencyResetInProgress = false;
      digitalWrite(LED_ALERT, LOW);
      digitalWrite(LED_SYSTEM_ONLINE, HIGH);
      Serial.println("✅ EMERGENCY STOP RESET via MQTT command!");
    } else {
      Serial.println("ℹ️ Emergency reset command received but no emergency stop active");
    }
  }
  
  // Individual control commands
  else if (topicStr == "industrial/control/motor") {
    if (!emergencyStop && systemEnabled) {
      bool state = (message == "on" || message == "1" || message == "true");
      digitalWrite(MOTOR_RELAY_PIN, state);
      controls.motorRunning = state;
      Serial.print("✅ Motor control: ");
      Serial.println(state ? "ON" : "OFF");
    } else {
      Serial.println("❌ Motor control BLOCKED - system not enabled or emergency stop");
    }
  }
  else if (topicStr == "industrial/control/pump") {
    if (!emergencyStop && systemEnabled) {
      bool state = (message == "on" || message == "1" || message == "true");
      digitalWrite(PUMP_RELAY_PIN, state);
      controls.pumpActive = state;
      Serial.print("✅ Pump control: ");
      Serial.println(state ? "ON" : "OFF");
    } else {
      Serial.println("❌ Pump control BLOCKED - system not enabled or emergency stop");
    }
  }
  else if (topicStr == "industrial/control/valve") {
    if (!emergencyStop && systemEnabled) {
      bool state = (message == "on" || message == "1" || message == "true");
      digitalWrite(VALVE_RELAY_PIN, state);
      controls.valveOpen = state;
      Serial.print("✅ Valve control: ");
      Serial.println(state ? "OPEN" : "CLOSED");
    } else {
      Serial.println("❌ Valve control BLOCKED - system not enabled or emergency stop");
    }
  }
  else if (topicStr == "industrial/control/cooling") {
    if (!emergencyStop && systemEnabled) {
      bool state = (message == "on" || message == "1" || message == "true");
      digitalWrite(COOLING_RELAY_PIN, state);
      controls.coolingActive = state;
      Serial.print("✅ Cooling control: ");
      Serial.println(state ? "ON" : "OFF");
    } else {
      Serial.println("❌ Cooling control BLOCKED - system not enabled or emergency stop");
    }
  }
  
  // LED Control Commands
  else if (topicStr == "industrial/led/red") {
    if (message == "auto") {
      ledRedManualOverride = false;
    } else {
      bool state = (message == "on" || message == "1" || message == "true");
      ledRedManualOverride = true;
      ledRedManualState = state;
      digitalWrite(LED_ALERT, state);
    }
    // Print updated LED status
    Serial.print("🚦 LED STATUS: Red LED: ");
    Serial.print(ledRedManualOverride ? (ledRedManualState ? "ON" : "OFF") : 
                (digitalRead(LED_ALERT) ? "ON" : "OFF"));
    Serial.print(" | Blue LED: ");
    Serial.println(ledBlueManualOverride ? (ledBlueManualState ? "ON" : "OFF") : 
                   (digitalRead(LED_PROCESS_ACTIVE) ? "ON" : "OFF"));
  }
  else if (topicStr == "industrial/led/blue") {
    if (message == "auto") {
      ledBlueManualOverride = false;
    } else {
      bool state = (message == "on" || message == "1" || message == "true");
      ledBlueManualOverride = true;
      ledBlueManualState = state;
      digitalWrite(LED_PROCESS_ACTIVE, state);
    }
    // Print updated LED status
    Serial.print("🚦 LED STATUS: Red LED: ");
    Serial.print(ledRedManualOverride ? (ledRedManualState ? "ON" : "OFF") : 
                (digitalRead(LED_ALERT) ? "ON" : "OFF"));
    Serial.print(" | Blue LED: ");
    Serial.println(ledBlueManualOverride ? (ledBlueManualState ? "ON" : "OFF") : 
                   (digitalRead(LED_PROCESS_ACTIVE) ? "ON" : "OFF"));
  }
  else if (topicStr == "industrial/led/green") {
    Serial.println("ℹ️ Green LED is always AUTOMATIC - shows system online status");
  }
  else if (topicStr == "industrial/led/yellow") {
    Serial.println("ℹ️ Yellow LED is always AUTOMATIC - shows power status");
  }
  
  // Sensor Configuration Commands
  else if (topicStr == "industrial/config/motion_enabled") {
    motionSensorEnabled = (message == "on" || message == "1" || message == "true");
    Serial.print("✅ Motion detected? : ");
    Serial.println(motionSensorEnabled ? "YES" : "NO");
    client.publish("industrial/config/motion_enabled/status", 
                   motionSensorEnabled ? "YES" : "NO", true);
  }
  else if (topicStr == "industrial/config/light_enabled") {
    lightSensorEnabled = (message == "on" || message == "1" || message == "true");
    Serial.print("✅ Light detected? : ");
    Serial.println(lightSensorEnabled ? "YES" : "NO");
    client.publish("industrial/config/light_enabled/status", 
                   lightSensorEnabled ? "YES" : "NO", true);
  }
  else {
    Serial.print("⚠️ UNKNOWN MQTT topic: ");
    Serial.println(topicStr);
  }
  
  Serial.println(String('=', 35));
}

void readSensors() {
  // Read DHT22 (Environment)
  float temp = dht.readTemperature();
  float humidity = dht.readHumidity();
  if (!isnan(temp)) sensors.environmentTemp = temp;
  if (!isnan(humidity)) sensors.environmentHumidity = humidity;
  
  // Read analog sensors (potentiometers simulate real sensors)
  int pressureRaw = analogRead(PRESSURE_PIN);
  sensors.pressure = map(pressureRaw, 0, 4095, PRESSURE_MIN * 10, PRESSURE_MAX * 10) / 10.0;
  
  int flowRaw = analogRead(FLOW_PIN);
  sensors.flowRate = map(flowRaw, 0, 4095, FLOW_MIN, FLOW_MAX);
  
  int vibrationRaw = analogRead(VIBRATION_PIN);
  sensors.vibration = map(vibrationRaw, 0, 4095, 0, 100);
  
  // Read light sensor (only if enabled)
  if (lightSensorEnabled) {
    int lightRaw = analogRead(LIGHT_PIN);
    sensors.lightLevel = map(lightRaw, 0, 4095, 0, 100);
  }
  
  // Read motion sensor (only if enabled)
  if (motionSensorEnabled) {
    sensors.motionDetected = digitalRead(MOTION_PIN);
  } else {
    sensors.motionDetected = false;
  }
}

void checkButtons() {
  static bool lastEmergencyState = HIGH;
  static bool lastStartState = HIGH;
  
  bool currentEmergencyState = digitalRead(BUTTON_EMERGENCY);
  bool currentStartState = digitalRead(BUTTON_START);
  
  // Emergency stop button pressed
  if (lastEmergencyState == HIGH && currentEmergencyState == LOW) {
    emergencyStop = true;
    currentState = EMERGENCY_STOP;
    stopAllOutputs();
    Serial.println("🚨 EMERGENCY STOP activated via button!");
  }
  
  // UPDATED: Emergency Reset Logic - Now works with single press too!
  if (emergencyStop) {
    if (currentStartState == LOW && lastStartState == HIGH) {
      // Single press - immediate reset
      emergencyStop = false;
      systemEnabled = false;
      currentState = SYSTEM_IDLE;
      emergencyResetInProgress = false;
      if (!ledRedManualOverride) {
        digitalWrite(LED_ALERT, LOW);
      }
      digitalWrite(LED_SYSTEM_ONLINE, HIGH);
      Serial.println("✅ EMERGENCY STOP RESET via button press!");
    }
    
    // Keep the old hold-for-3-seconds method as backup
    if (currentStartState == LOW) {
      if (!emergencyResetInProgress) {
        emergencyResetStartTime = millis();
        emergencyResetInProgress = true;
      } else {
        // Check if button has been held long enough
        if (millis() - emergencyResetStartTime >= EMERGENCY_RESET_TIME) {
          emergencyStop = false;
          systemEnabled = false;
          currentState = SYSTEM_IDLE;
          emergencyResetInProgress = false;
          if (!ledRedManualOverride) {
            digitalWrite(LED_ALERT, LOW);
          }
          digitalWrite(LED_SYSTEM_ONLINE, HIGH);
          Serial.println("✅ EMERGENCY STOP RESET via button hold!");
        }
      }
    } else {
      // Button released, cancel hold reset
      if (emergencyResetInProgress) {
        emergencyResetInProgress = false;
      }
    }
  } else {
    // Normal start button operation
    if (lastStartState == HIGH && currentStartState == LOW) {
      systemEnabled = true;
      currentState = SYSTEM_STARTUP;
      Serial.println("✅ System STARTED via button!");
    }
  }
  
  lastEmergencyState = currentEmergencyState;
  lastStartState = currentStartState;
}

void executeStateMachine() {
  switch (currentState) {
    case SYSTEM_IDLE:
      handleIdleState();
      break;
      
    case SYSTEM_STARTUP:
      handleStartupState();
      break;
      
    case SYSTEM_ACTIVE:  // Combined running and monitoring
      handleActiveState();
      break;
      
    case SYSTEM_ALERT:
      handleAlertState();
      break;
      
    case EMERGENCY_STOP:
      handleEmergencyState();
      break;
  }
}

void handleIdleState() {
  stopAllOutputs();
  if (!ledBlueManualOverride) {
    digitalWrite(LED_PROCESS_ACTIVE, LOW);
  }
  processActive = false;
  
  if (systemEnabled && !emergencyStop) {
    currentState = SYSTEM_STARTUP;
  }
}

void handleStartupState() {
  static unsigned long startupTime = 0;
  if (startupTime == 0) startupTime = millis();
  
  // Startup sequence - activate systems gradually (non-blocking)
  if (millis() - startupTime > 2000) {
    if (!ledBlueManualOverride) {
      digitalWrite(LED_PROCESS_ACTIVE, HIGH);
    }
    processActive = true;
    currentState = SYSTEM_ACTIVE;
    startupTime = 0;
  }
}

void handleActiveState() {
  // Check for emergency or system disable first
  if (!systemEnabled || emergencyStop) {
    currentState = SYSTEM_IDLE;
    return;
  }
  
  // Check for alert conditions
  if (temperatureAlert || pressureAlert || vibrationAlert) {
    currentState = SYSTEM_ALERT;
    return;
  }
  
  // Automatic process control based on sensor readings
  
  // Temperature control
  if (sensors.environmentTemp > TEMP_HIGH_THRESHOLD && !controls.coolingActive) {
    digitalWrite(COOLING_RELAY_PIN, HIGH);
    controls.coolingActive = true;
  } else if (sensors.environmentTemp < (TEMP_HIGH_THRESHOLD - 2) && controls.coolingActive) {
    digitalWrite(COOLING_RELAY_PIN, LOW);
    controls.coolingActive = false;
  }
  
  // Flow control
  if (sensors.pressure > 0.5 && !controls.valveOpen) {
    digitalWrite(VALVE_RELAY_PIN, HIGH);
    controls.valveOpen = true;
  }
  
  // Pump control based on flow demand
  if (sensors.flowRate > 50 && !controls.pumpActive) {
    digitalWrite(PUMP_RELAY_PIN, HIGH);
    controls.pumpActive = true;
  }
  
  // Continue monitoring in this state
}

void handleAlertState() {
  // Handle alert conditions
  
  // Flash alert LED (non-blocking)
  if (!ledRedManualOverride) {
    static unsigned long lastFlash = 0;
    if (millis() - lastFlash > 500) {
      digitalWrite(LED_ALERT, !digitalRead(LED_ALERT));
      lastFlash = millis();
    }
  }
  
  // Sound buzzer periodically (non-blocking)
  static unsigned long lastBuzzer = 0;
  if (millis() - lastBuzzer > 3000) {
    tone(BUZZER_PIN, 1500, 300);
    lastBuzzer = millis();
  }
  
  // Check if alerts cleared
  if (!temperatureAlert && !pressureAlert && !vibrationAlert) {
    if (!ledRedManualOverride) {
      digitalWrite(LED_ALERT, LOW);
    }
    currentState = SYSTEM_ACTIVE;
  }
  
  if (!systemEnabled || emergencyStop) {
    currentState = SYSTEM_IDLE;
  }
}

void handleEmergencyState() {
  stopAllOutputs();
  if (!ledRedManualOverride) {
    digitalWrite(LED_ALERT, HIGH);
  }
  if (!ledBlueManualOverride) {
    digitalWrite(LED_PROCESS_ACTIVE, LOW);
  }
  digitalWrite(LED_SYSTEM_ONLINE, LOW);
  
  // Continuous alarm for first 5 seconds (non-blocking)
  static bool alarmSounded = false;
  static unsigned long alarmStart = 0;
  
  if (!alarmSounded) {
    tone(BUZZER_PIN, 2500);
    alarmSounded = true;
    alarmStart = millis();
  }
  
  if (millis() - alarmStart > 5000) {
    noTone(BUZZER_PIN);
  }
  
  // Reset handled in checkButtons() function
  if (!emergencyStop) {
    alarmSounded = false;
    if (!ledRedManualOverride) {
      digitalWrite(LED_ALERT, LOW);
    }
    digitalWrite(LED_SYSTEM_ONLINE, HIGH);
  }
}

void performAlertChecks() {
  // Temperature alerts
  temperatureAlert = (sensors.environmentTemp > (TEMP_HIGH_THRESHOLD + 10));
  
  // Pressure alerts
  pressureAlert = (sensors.pressure > (PRESSURE_MAX - 1)) || 
                  (sensors.pressure < 0);
  
  // Vibration alerts
  vibrationAlert = (sensors.vibration > VIBRATION_THRESHOLD);
  
  // Publish critical alerts (non-blocking)
  if ((temperatureAlert || pressureAlert || vibrationAlert) && 
      (millis() - lastAlertTime > 15000)) {
    if (mqttConnected) {
      if (temperatureAlert) {
        client.publish("industrial/alert/temperature", "CRITICAL");
      }
      if (pressureAlert) {
        client.publish("industrial/alert/pressure", "CRITICAL");
      }
      if (vibrationAlert) {
        client.publish("industrial/alert/vibration", "WARNING");
      }
    }
    lastAlertTime = millis();
  }
}

void updateStatusLEDs() {
  // System online LED (Green) - Always automatic control
  if (wifiConnected && mqttConnected && !emergencyStop) {
    digitalWrite(LED_SYSTEM_ONLINE, HIGH);
  } else {
    static unsigned long lastBlink = 0;
    if (millis() - lastBlink > 1000) {
      digitalWrite(LED_SYSTEM_ONLINE, !digitalRead(LED_SYSTEM_ONLINE));
      lastBlink = millis();
    }
  }
  
  // Alert LED (Red) - Check for manual override first
  if (ledRedManualOverride) {
    digitalWrite(LED_ALERT, ledRedManualState);
  } else {
    // Automatic control based on system state
    if (currentState == SYSTEM_ALERT) {
      // Flashing handled in handleAlertState() function
    } else if (currentState == EMERGENCY_STOP) {
      digitalWrite(LED_ALERT, HIGH);
    } else {
      digitalWrite(LED_ALERT, LOW);
    }
  }
  
  // Process active LED (Blue) - Check for manual override first
  if (ledBlueManualOverride) {
    digitalWrite(LED_PROCESS_ACTIVE, ledBlueManualState);
  } else {
    // Automatic control based on process state
    digitalWrite(LED_PROCESS_ACTIVE, processActive && !emergencyStop);
  }
  
  // Power status LED (Yellow) - Always automatic control
  digitalWrite(LED_POWER_STATUS, !emergencyStop);
}

void publishSensorData() {
  if (!mqttConnected) return;
  
  // Create comprehensive telemetry JSON
  StaticJsonDocument<1000> doc;
  
  // Environmental data
  doc["environment"]["temperature"] = sensors.environmentTemp;
  doc["environment"]["humidity"] = sensors.environmentHumidity;
  doc["environment"]["light_level"] = sensors.lightLevel;
  doc["environment"]["motion_detected"] = sensors.motionDetected;
  
  // Process data
  doc["process"]["pressure"] = sensors.pressure;
  doc["process"]["flow_rate"] = sensors.flowRate;
  doc["process"]["vibration"] = sensors.vibration;
  
  // Control states
  doc["controls"]["motor_running"] = controls.motorRunning;
  doc["controls"]["pump_active"] = controls.pumpActive;
  doc["controls"]["valve_open"] = controls.valveOpen;
  doc["controls"]["cooling_active"] = controls.coolingActive;
  
  // System status
  doc["system"]["state"] = stateToString(currentState);
  doc["system"]["enabled"] = systemEnabled;
  doc["system"]["emergency_stop"] = emergencyStop;
  doc["system"]["process_active"] = processActive;
  
  // Alerts
  doc["alerts"]["temperature"] = temperatureAlert;
  doc["alerts"]["pressure"] = pressureAlert;
  doc["alerts"]["vibration"] = vibrationAlert;
  
  // Configuration states
  doc["config"]["motion_sensor_enabled"] = motionSensorEnabled;
  doc["config"]["light_sensor_enabled"] = lightSensorEnabled;
  
  // LED states
  doc["leds"]["red_manual"] = ledRedManualOverride;
  doc["leds"]["blue_manual"] = ledBlueManualOverride;
  doc["leds"]["red_state"] = digitalRead(LED_ALERT);
  doc["leds"]["blue_state"] = digitalRead(LED_PROCESS_ACTIVE);
  doc["leds"]["green_state"] = digitalRead(LED_SYSTEM_ONLINE);
  doc["leds"]["yellow_state"] = digitalRead(LED_POWER_STATUS);
  
  doc["timestamp"] = millis();
  
  char buffer[1000];
  serializeJson(doc, buffer);
  
  client.publish("industrial/telemetry", buffer);
  
  // Individual sensor topics for easier monitoring
  client.publish("industrial/sensors/env_temp", String(sensors.environmentTemp).c_str());
  client.publish("industrial/sensors/pressure", String(sensors.pressure).c_str());
  client.publish("industrial/sensors/flow_rate", String(sensors.flowRate).c_str());
  client.publish("industrial/sensors/vibration", String(sensors.vibration).c_str());
  client.publish("industrial/sensors/light_level", String(sensors.lightLevel).c_str());
  client.publish("industrial/sensors/motion_detected", sensors.motionDetected ? "true" : "false");
}

void publishConfigurationStates() {
  if (!mqttConnected) return;
  
  // Publish current configuration states
  client.publish("industrial/config/motion_enabled/status", 
                 motionSensorEnabled ? "enabled" : "disabled", true);
  client.publish("industrial/config/light_enabled/status", 
                 lightSensorEnabled ? "enabled" : "disabled", true);
  
  // Publish LED override states
  client.publish("industrial/led/red/status", 
                 ledRedManualOverride ? (ledRedManualState ? "manual_on" : "manual_off") : "auto", true);
  client.publish("industrial/led/blue/status", 
                 ledBlueManualOverride ? (ledBlueManualState ? "manual_on" : "manual_off") : "auto", true);
}

void stopAllOutputs() {
  digitalWrite(MOTOR_RELAY_PIN, LOW);
  digitalWrite(PUMP_RELAY_PIN, LOW);
  digitalWrite(VALVE_RELAY_PIN, LOW);
  digitalWrite(COOLING_RELAY_PIN, LOW);
  noTone(BUZZER_PIN);
  
  controls.motorRunning = false;
  controls.pumpActive = false;
  controls.valveOpen = false;
  controls.coolingActive = false;
}

void performSystemCheck() {
  Serial.println("🔍 Performing comprehensive system check...");
  
  updateLcdLine(0, "System Check");
  updateLcdLine(1, "Initializing...");
  
  // Test all LEDs
  Serial.println("   Testing LED indicators...");
  digitalWrite(LED_SYSTEM_ONLINE, HIGH);
  digitalWrite(LED_ALERT, HIGH);
  digitalWrite(LED_PROCESS_ACTIVE, HIGH);
  digitalWrite(LED_POWER_STATUS, HIGH);
  delay(1000);
  
  digitalWrite(LED_SYSTEM_ONLINE, LOW);
  digitalWrite(LED_ALERT, LOW);
  digitalWrite(LED_PROCESS_ACTIVE, LOW);
  // Keep power LED on
  
  // Test buzzer
  Serial.println("   Testing alarm system...");
  tone(BUZZER_PIN, 1000, 500);
  delay(600);
  
  updateLcdLine(1, "Check Complete");
  delay(1500);
  
  Serial.println("✅ System hardware check completed successfully");
}

void updateLcdLine(int line, String content) {
  // Pad content to 16 characters
  while (content.length() < 16) {
    content += " ";
  }
  if (content.length() > 16) {
    content = content.substring(0, 16);
  }
  
  if (lcdLine[line] != content) {
    lcdLine[line] = content;
    lcd.setCursor(0, line);
    lcd.print(content);
  }
}

void updateLcdDisplay() {
  // Rotate display information every cycle
  static int displayMode = 0;
  static unsigned long lastModeChange = 0;
  
  if (millis() - lastModeChange > 10000) { // Changed to 10 seconds (slower)
    displayMode = (displayMode + 1) % 5;
    lastModeChange = millis();
  }
  
  switch (displayMode) {
    case 0:
      // System status
      updateLcdLine(0, "Status:" + String(stateToString(currentState)));
      updateLcdLine(1, "T:" + String(sensors.environmentTemp, 1) + "C P:" + String(sensors.pressure, 1) + "bar");
      break;
      
    case 1:
      // Environmental data
      updateLcdLine(0, "Env: " + String(sensors.environmentTemp, 1) + "C " + String(sensors.environmentHumidity, 0) + "%");
      updateLcdLine(1, "Light:" + String(sensors.lightLevel) + "% Flow:" + String(sensors.flowRate, 0));
      break;
      
    case 2:
      {
        // Control status
        String line1 = "M:";
        line1 += controls.motorRunning ? "ON " : "OFF ";
        line1 += "P:";
        line1 += controls.pumpActive ? "ON" : "OFF";
        updateLcdLine(0, line1);
        
        String line2 = "V:";
        line2 += controls.valveOpen ? "OPEN " : "SHUT ";
        line2 += "C:";
        line2 += controls.coolingActive ? "ON" : "OFF";
        updateLcdLine(1, line2);
        break;
      }
      
    case 3:
      // Connection and alerts
      if (temperatureAlert || pressureAlert || vibrationAlert) {
        updateLcdLine(0, "*** ALERTS ***");
        String alerts = "";
        if (temperatureAlert) alerts += "TEMP ";
        if (pressureAlert) alerts += "PRESS ";
        if (vibrationAlert) alerts += "VIB";
        updateLcdLine(1, alerts);
      } else if (!wifiConnected) {
        updateLcdLine(0, "WiFi Disconnected");
        updateLcdLine(1, "Check network");
      } else if (!mqttConnected) {
        updateLcdLine(0, "MQTT Disconnected");
        updateLcdLine(1, "Broker offline");
      } else {
        updateLcdLine(0, "All Systems OK");
        updateLcdLine(1, "Vibration:" + String(sensors.vibration, 0) + "%");
      }
      break;
      
    case 4:
      // Sensor configuration status
      {
        String line1 = "Sensors: ";
        line1 += motionSensorEnabled ? "M" : "-";
        line1 += lightSensorEnabled ? "L " : "- ";
        if (emergencyResetInProgress) {
          float remaining = (EMERGENCY_RESET_TIME - (millis() - emergencyResetStartTime)) / 1000.0;
          line1 += String(remaining, 1) + "s";
        }
        updateLcdLine(0, line1);
        
        String line2 = "LEDs: ";
        line2 += ledRedManualOverride ? "R" : "-";
        line2 += ledBlueManualOverride ? "B " : "- ";
        line2 += "Motion:" + String(sensors.motionDetected ? "Y" : "N");
        updateLcdLine(1, line2);
        break;
      }
  }
}

const char* stateToString(SystemState state) {
  switch (state) {
    case SYSTEM_IDLE: return "IDLE";
    case SYSTEM_STARTUP: return "STARTUP";
    case SYSTEM_ACTIVE: return "ACTIVE";
    case SYSTEM_ALERT: return "ALERT";
    case EMERGENCY_STOP: return "E-STOP";
    default: return "UNKNOWN";
  }
}