/*
 * ESP32 Containment Unit - Build 1.0
 * Improved version with state machine and better organization
 */

#include <Arduino.h>

// ===== PIN DEFINITIONS =====
// Input Pins
const int POWER_BUTTON_PIN = 13;      // Power Button
const int LIMIT_SWITCH_PIN = 34;      // Door position sensor
const int GREEN_SWITCH_PIN = 32;      // Ready Switch
const int RED_SWITCH_PIN = 33;        // Containment Lock
const int ORANGE_SWITCH_PIN = 4;      // Final Lock
const int LEVER_SWITCH_PIN = 15;      // Final Step Lever

// Ultrasonic Sensor Pins
const int TRIG_PIN = 5;               // Ultrasonic Trigger
const int ECHO_PIN = 18;              // Ultrasonic Echo

// LED Pins
const int GREEN_LED_PIN = 22;         // System Ready Indicator
const int RED_LED_PIN = 23;           // Containment Warning

// Relay Pins
const int RELAY_GREEN_PIN = 19;       // Green Switch Relay
const int RELAY_RED_PIN = 21;         // Red Switch Relay
const int RELAY_ORANGE_PIN = 25;      // Orange Switch Relay
const int RELAY_SMOKE_PIN = 26;       // Smoke Effect Relay
const int RELAY_LOCK_PIN = 27;        // Containment Lock Relay

// ===== CONFIGURATION CONSTANTS =====
const int TRAP_DISTANCE_THRESHOLD_CM = 20;    // Distance to detect trap insertion
const int ULTRASONIC_TIMEOUT_US = 30000;      // Ultrasonic sensor timeout
const unsigned long SMOKE_DURATION_MS = 2000; // How long smoke effect runs
const unsigned long DEBOUNCE_DELAY_MS = 50;   // Button debounce time
const unsigned long SENSOR_READ_INTERVAL_MS = 100; // How often to read ultrasonic
const unsigned long SERIAL_BAUD_RATE = 115200;

// Relay States (Active Low for most relay modules)
const int RELAY_ON = LOW;
const int RELAY_OFF = HIGH;

// ===== SYSTEM STATES =====
enum SystemState {
  STATE_OFF,                    // System powered off
  STATE_READY,                  // System on, waiting for door open
  STATE_DOOR_OPEN,             // Door open, waiting for trap
  STATE_TRAP_DETECTED,         // Trap detected, waiting for door close
  STATE_DOOR_CLOSED,           // Door closed, waiting for red button
  STATE_RED_PRESSED,           // Red pressed, waiting for orange button
  STATE_ORANGE_PRESSED,        // Orange pressed, waiting for lever down
  STATE_LEVER_DOWN,            // Lever down, smoke activated
  STATE_PROCESSING,            // Processing complete, waiting for lever up
  STATE_ERROR                  // Error state
};

// ===== GLOBAL VARIABLES =====
SystemState currentState = STATE_OFF;
SystemState previousState = STATE_OFF;

// Timing variables for non-blocking operations
unsigned long lastSensorRead = 0;
unsigned long smokeStartTime = 0;
unsigned long lastDebounceTime = 0;

// Button state tracking for debouncing
struct ButtonState {
  bool currentState = HIGH;
  bool lastState = HIGH;
  bool stableState = HIGH;
  unsigned long lastChangeTime = 0;
};

ButtonState powerButton, greenSwitch, redSwitch, orangeSwitch, leverSwitch, limitSwitch;

// Sensor data
int currentDistance = 0;
bool trapDetected = false;

// ===== SETUP FUNCTION =====
void setup() {
  Serial.begin(SERIAL_BAUD_RATE);
  delay(100); // Allow serial to stabilize
  
  Serial.println("🚀 ESP32 Containment Unit Starting...");
  
  setupPins();
  resetSystem();
  
  Serial.println("✅ System Initialized - Press Power Button to Start");
}

// ===== MAIN LOOP =====
void loop() {
  updateButtonStates();
  updateSensors();
  handleStateMachine();
  updateOutputs();
  
  // Small delay to prevent excessive CPU usage
  delay(10);
}

// ===== PIN SETUP =====
void setupPins() {
  // Setup input pins with internal pull-ups
  pinMode(POWER_BUTTON_PIN, INPUT_PULLUP);
  pinMode(GREEN_SWITCH_PIN, INPUT_PULLUP);
  pinMode(RED_SWITCH_PIN, INPUT_PULLUP);
  pinMode(ORANGE_SWITCH_PIN, INPUT_PULLUP);
  pinMode(LEVER_SWITCH_PIN, INPUT_PULLUP);
  pinMode(LIMIT_SWITCH_PIN, INPUT_PULLUP);
  
  // Setup ultrasonic sensor pins
  pinMode(TRIG_PIN, OUTPUT);
  pinMode(ECHO_PIN, INPUT);
  digitalWrite(TRIG_PIN, LOW);
  
  // Setup LED pins
  pinMode(GREEN_LED_PIN, OUTPUT);
  pinMode(RED_LED_PIN, OUTPUT);
  
  // Setup relay pins
  pinMode(RELAY_GREEN_PIN, OUTPUT);
  pinMode(RELAY_RED_PIN, OUTPUT);
  pinMode(RELAY_ORANGE_PIN, OUTPUT);
  pinMode(RELAY_SMOKE_PIN, OUTPUT);
  pinMode(RELAY_LOCK_PIN, OUTPUT);
  
  Serial.println("📌 Pins configured successfully");
}

// ===== BUTTON STATE MANAGEMENT =====
void updateButtonStates() {
  updateButtonState(powerButton, POWER_BUTTON_PIN);
  updateButtonState(greenSwitch, GREEN_SWITCH_PIN);
  updateButtonState(redSwitch, RED_SWITCH_PIN);
  updateButtonState(orangeSwitch, ORANGE_SWITCH_PIN);
  updateButtonState(leverSwitch, LEVER_SWITCH_PIN);
  updateButtonState(limitSwitch, LIMIT_SWITCH_PIN);
}

void updateButtonState(ButtonState &button, int pin) {
  bool reading = digitalRead(pin);
  unsigned long currentTime = millis();
  
  // If the switch changed, due to noise or pressing
  if (reading != button.lastState) {
    button.lastChangeTime = currentTime;
  }
  
  // If the reading has been stable for the debounce delay
  if ((currentTime - button.lastChangeTime) > DEBOUNCE_DELAY_MS) {
    if (reading != button.stableState) {
      button.stableState = reading;
    }
  }
  
  button.lastState = reading;
}

bool isButtonPressed(const ButtonState &button) {
  return (button.stableState == LOW);
}

bool isButtonJustPressed(ButtonState &button) {
  if (button.stableState == LOW && button.currentState == HIGH) {
    button.currentState = LOW;
    return true;
  } else if (button.stableState == HIGH) {
    button.currentState = HIGH;
  }
  return false;
}

bool isButtonJustReleased(ButtonState &button) {
  if (button.stableState == HIGH && button.currentState == LOW) {
    button.currentState = HIGH;
    return true;
  }
  return false;
}

// ===== SENSOR MANAGEMENT =====
void updateSensors() {
  unsigned long currentTime = millis();
  
  // Read ultrasonic sensor at intervals to avoid blocking
  if (currentTime - lastSensorRead >= SENSOR_READ_INTERVAL_MS) {
    lastSensorRead = currentTime;
    currentDistance = measureDistance();
    trapDetected = (currentDistance > 0 && currentDistance < TRAP_DISTANCE_THRESHOLD_CM);
  }
}

int measureDistance() {
  // Send ultrasonic pulse
  digitalWrite(TRIG_PIN, LOW);
  delayMicroseconds(2);
  digitalWrite(TRIG_PIN, HIGH);
  delayMicroseconds(10);
  digitalWrite(TRIG_PIN, LOW);
  
  // Measure echo duration with timeout
  long duration = pulseIn(ECHO_PIN, HIGH, ULTRASONIC_TIMEOUT_US);
  
  if (duration == 0) {
    return -1; // Timeout or no echo
  }
  
  // Convert to centimeters
  int distance = duration * 0.034 / 2;
  return distance;
}

// ===== STATE MACHINE =====
void handleStateMachine() {
  SystemState newState = currentState;
  
  switch (currentState) {
    case STATE_OFF:
      newState = handleStateOff();
      break;
      
    case STATE_READY:
      newState = handleStateReady();
      break;
      
    case STATE_DOOR_OPEN:
      newState = handleStateDoorOpen();
      break;
      
    case STATE_TRAP_DETECTED:
      newState = handleStateTrapDetected();
      break;
      
    case STATE_DOOR_CLOSED:
      newState = handleStateDoorClosed();
      break;
      
    case STATE_RED_PRESSED:
      newState = handleStateRedPressed();
      break;
      
    case STATE_ORANGE_PRESSED:
      newState = handleStateOrangePressed();
      break;
      
    case STATE_LEVER_DOWN:
      newState = handleStateLeverDown();
      break;
      
    case STATE_PROCESSING:
      newState = handleStateProcessing();
      break;
      
    case STATE_ERROR:
      newState = handleStateError();
      break;
  }
  
  // Handle state transitions
  if (newState != currentState) {
    changeState(newState);
  }
  
  // Global power button check (can turn off from any state)
  if (isButtonJustPressed(powerButton) && currentState != STATE_OFF) {
    changeState(STATE_OFF);
  }
}

// ===== STATE HANDLERS =====
SystemState handleStateOff() {
  if (isButtonJustPressed(powerButton)) {
    Serial.println("✅ SYSTEM POWERED ON - Ready Mode");
    return STATE_READY;
  }
  return STATE_OFF;
}

SystemState handleStateReady() {
  if (!isButtonPressed(limitSwitch)) { // Door opened (limit switch released)
    Serial.println("🚪 DOOR OPENED - Waiting for Trap Insertion");
    return STATE_DOOR_OPEN;
  }
  return STATE_READY;
}

SystemState handleStateDoorOpen() {
  if (trapDetected) {
    Serial.println("👻 TRAP DETECTED - Close Door to Continue");
    Serial.printf("📏 Distance: %d cm\n", currentDistance);
    return STATE_TRAP_DETECTED;
  }
  
  if (isButtonPressed(limitSwitch)) { // Door closed without trap
    Serial.println("⚠️  Door closed without trap - Returning to Ready");
    return STATE_READY;
  }
  
  return STATE_DOOR_OPEN;
}

SystemState handleStateTrapDetected() {
  if (isButtonPressed(limitSwitch)) { // Door closed with trap inside
    Serial.println("🚪 DOOR CLOSED WITH TRAP - Press Red Button to Proceed");
    return STATE_DOOR_CLOSED;
  }
  
  if (!trapDetected) { // Trap removed
    Serial.println("❌ Trap removed - Waiting for trap insertion");
    return STATE_DOOR_OPEN;
  }
  
  return STATE_TRAP_DETECTED;
}

SystemState handleStateDoorClosed() {
  if (isButtonJustPressed(redSwitch)) {
    Serial.println("🔴 RED BUTTON PRESSED - Press Orange Button");
    return STATE_RED_PRESSED;
  }
  
  if (!isButtonPressed(limitSwitch)) { // Door reopened
    Serial.println("🚪 Door reopened - Returning to door open state");
    return STATE_DOOR_OPEN;
  }
  
  return STATE_DOOR_CLOSED;
}

SystemState handleStateRedPressed() {
  if (isButtonJustPressed(orangeSwitch)) {
    Serial.println("🟠 ORANGE BUTTON PRESSED - Throw Lever Down");
    return STATE_ORANGE_PRESSED;
  }
  return STATE_RED_PRESSED;
}

SystemState handleStateOrangePressed() {
  if (isButtonJustPressed(leverSwitch)) {
    Serial.println("🔄 LEVER THROWN DOWN - Activating Smoke Effect!");
    smokeStartTime = millis();
    return STATE_LEVER_DOWN;
  }
  return STATE_ORANGE_PRESSED;
}

SystemState handleStateLeverDown() {
  unsigned long currentTime = millis();
  
  // Check if smoke duration has elapsed
  if (currentTime - smokeStartTime >= SMOKE_DURATION_MS) {
    Serial.println("💨 Smoke effect complete - Throw Lever Up to Reset");
    return STATE_PROCESSING;
  }
  
  return STATE_LEVER_DOWN;
}

SystemState handleStateProcessing() {
  if (isButtonJustReleased(leverSwitch)) {
    Serial.println("🔼 LEVER THROWN UP - CONTAINMENT COMPLETE!");
    Serial.println("🔄 System Ready for Next Operation");
    return STATE_READY;
  }
  return STATE_PROCESSING;
}

SystemState handleStateError() {
  // Error recovery - return to ready state if power cycled
  if (isButtonJustPressed(powerButton)) {
    Serial.println("🔧 Error cleared - Returning to Ready");
    return STATE_READY;
  }
  return STATE_ERROR;
}

// ===== STATE MANAGEMENT =====
void changeState(SystemState newState) {
  previousState = currentState;
  currentState = newState;
  
  Serial.printf("🔄 State: %s → %s\n", 
    getStateName(previousState).c_str(), 
    getStateName(newState).c_str());
  
  // Handle state entry actions
  onStateEnter(newState);
}

void onStateEnter(SystemState state) {
  switch (state) {
    case STATE_OFF:
      resetSystem();
      break;
      
    case STATE_READY:
      // System ready - show green status
      break;
      
    case STATE_DOOR_OPEN:
      // Door is open - dim green, prepare for trap
      break;
      
    case STATE_TRAP_DETECTED:
      // Trap detected - show warning
      break;
      
    case STATE_DOOR_CLOSED:
      // Door closed with trap - ready for sequence
      break;
      
    case STATE_RED_PRESSED:
      // Red button activated
      break;
      
    case STATE_ORANGE_PRESSED:
      // Orange button activated
      break;
      
    case STATE_LEVER_DOWN:
      // Start smoke effect
      smokeStartTime = millis();
      break;
      
    case STATE_PROCESSING:
      // Processing complete
      break;
      
    case STATE_ERROR:
      Serial.println("❌ ERROR STATE ENTERED");
      break;
  }
}

// ===== OUTPUT MANAGEMENT =====
void updateOutputs() {
  updateLEDs();
  updateRelays();
}

void updateLEDs() {
  bool greenLED = false;
  bool redLED = false;
  
  switch (currentState) {
    case STATE_OFF:
      // All LEDs off
      break;
      
    case STATE_READY:
      greenLED = true;
      break;
      
    case STATE_DOOR_OPEN:
      // Slow blink green
      greenLED = (millis() % 1000) < 500;
      break;
      
    case STATE_TRAP_DETECTED:
    case STATE_DOOR_CLOSED:
      redLED = true;
      break;
      
    case STATE_RED_PRESSED:
    case STATE_ORANGE_PRESSED:
      // Fast blink red
      redLED = (millis() % 200) < 100;
      break;
      
    case STATE_LEVER_DOWN:
      // Both LEDs on during processing
      greenLED = true;
      redLED = true;
      break;
      
    case STATE_PROCESSING:
      greenLED = true;
      break;
      
    case STATE_ERROR:
      // Fast blink both
      greenLED = redLED = (millis() % 100) < 50;
      break;
  }
  
  digitalWrite(GREEN_LED_PIN, greenLED);
  digitalWrite(RED_LED_PIN, redLED);
}

void updateRelays() {
  // Green switch relay (illumination)
  bool greenSwitchOn = (currentState == STATE_READY || currentState == STATE_PROCESSING);
  digitalWrite(RELAY_GREEN_PIN, greenSwitchOn ? RELAY_ON : RELAY_OFF);
  
  // Red switch relay (illumination)
  bool redSwitchOn = (currentState == STATE_RED_PRESSED || currentState == STATE_ORANGE_PRESSED);
  digitalWrite(RELAY_RED_PIN, redSwitchOn ? RELAY_ON : RELAY_OFF);
  
  // Orange switch relay (illumination)
  bool orangeSwitchOn = (currentState == STATE_ORANGE_PRESSED);
  digitalWrite(RELAY_ORANGE_PIN, orangeSwitchOn ? RELAY_ON : RELAY_OFF);
  
  // Smoke relay
  bool smokeOn = (currentState == STATE_LEVER_DOWN);
  digitalWrite(RELAY_SMOKE_PIN, smokeOn ? RELAY_ON : RELAY_OFF);
  
  // Lock relay (could be used for actual containment locking)
  bool lockOn = (currentState >= STATE_DOOR_CLOSED && currentState <= STATE_PROCESSING);
  digitalWrite(RELAY_LOCK_PIN, lockOn ? RELAY_ON : RELAY_OFF);
}

// ===== UTILITY FUNCTIONS =====
String getStateName(SystemState state) {
  switch (state) {
    case STATE_OFF: return "OFF";
    case STATE_READY: return "READY";
    case STATE_DOOR_OPEN: return "DOOR_OPEN";
    case STATE_TRAP_DETECTED: return "TRAP_DETECTED";
    case STATE_DOOR_CLOSED: return "DOOR_CLOSED";
    case STATE_RED_PRESSED: return "RED_PRESSED";
    case STATE_ORANGE_PRESSED: return "ORANGE_PRESSED";
    case STATE_LEVER_DOWN: return "LEVER_DOWN";
    case STATE_PROCESSING: return "PROCESSING";
    case STATE_ERROR: return "ERROR";
    default: return "UNKNOWN";
  }
}

void resetSystem() {
  // Turn off all outputs
  digitalWrite(GREEN_LED_PIN, LOW);
  digitalWrite(RED_LED_PIN, LOW);
  digitalWrite(RELAY_GREEN_PIN, RELAY_OFF);
  digitalWrite(RELAY_RED_PIN, RELAY_OFF);
  digitalWrite(RELAY_ORANGE_PIN, RELAY_OFF);
  digitalWrite(RELAY_SMOKE_PIN, RELAY_OFF);
  digitalWrite(RELAY_LOCK_PIN, RELAY_OFF);
  
  // Reset variables
  trapDetected = false;
  currentDistance = 0;
  
  Serial.println("🔄 System Reset Complete");
}

// ===== DEBUG FUNCTIONS =====
void printSystemStatus() {
  Serial.println("=== SYSTEM STATUS ===");
  Serial.printf("State: %s\n", getStateName(currentState).c_str());
  Serial.printf("Distance: %d cm\n", currentDistance);
  Serial.printf("Trap Detected: %s\n", trapDetected ? "YES" : "NO");
  Serial.printf("Door Open: %s\n", !isButtonPressed(limitSwitch) ? "YES" : "NO");
  Serial.printf("Power: %s\n", isButtonPressed(powerButton) ? "PRESSED" : "RELEASED");
  Serial.println("====================");
}