#include <Arduino.h>
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <Adafruit_NeoPixel.h>
#include <Keypad.h>

// Display Configuration
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// NeoPixel Configuration
#define DOOR_NEOPIXEL_PIN 6
#define DOOR_LEDS 16
Adafruit_NeoPixel doorRing(DOOR_LEDS, DOOR_NEOPIXEL_PIN, NEO_GRB + NEO_KHZ800);

// Pin Definitions - FIXED ASSIGNMENTS
#define BUZZER_PIN 9
#define SOLENOID_RELAY_PIN 2
#define LIGHT_RELAY_PIN 3
#define OUTSIDE_IR 4
#define INSIDE_IR 5
#define LIGHT_SENSOR A0
#define SMOKE_ANALOG A1
#define SMOKE_DIGITAL 11    // Changed from A2 (was analog-only issue)
#define SMOKE_SW 12         // Changed from A3 (was analog-only issue)

// Keypad Configuration - FIXED PIN ASSIGNMENTS
const byte ROWS = 4;
const byte COLS = 4;
char keys[ROWS][COLS] = {
  {'1', '2', '3', 'A'},
  {'4', '5', '6', 'B'},
  {'7', '8', '9', 'C'},
  {'*', '0', '#', 'D'}
};
// Changed from problematic A6, A7 which are analog-only on Uno
byte rowPins[ROWS] = {13, A2, A3, A4};
byte colPins[COLS] = {A5, 10, 8, 7};
Keypad keypad = Keypad(makeKeymap(keys), rowPins, colPins, ROWS, COLS);

// System Variables
const char CORRECT_PASSWORD[] = "1234";
char enteredPassword[5] = {0};
enum DoorState { MAIN_MENU, ENTER_PASSWORD, UNLOCKED, WRONG, LOCKED_OUT };
DoorState doorState = MAIN_MENU;
bool doorUnlocked = false;
unsigned long lockoutStartTime = 0;
int wrongAttempts = 0;
const int MAX_ATTEMPTS = 3;
unsigned long unlockTime = 0;
const unsigned long UNLOCK_DURATION = 8000UL;
const unsigned long LOCKOUT_DURATION = 60000UL;
int peopleCount = 0;
bool lightsOn = false;
bool smokeDetected = false;

// Lighting thresholds - FIXED: Now using constants properly
const long DARK_THRESHOLD_LUX = 10000;  // Below this = dark, lights needed
const int SMOKE_THRESHOLD = 400;  // Gas sensor analog threshold (adjusted for typical values)

bool fireAlarmHigh = true;
unsigned long fireAlarmTimer = 0;

// PIR detection variables
bool lastOutsidePIR = false;
bool lastInsidePIR = false;
unsigned long lastOutsideMotionTime = 0;
unsigned long lastInsideMotionTime = 0;
const unsigned long MOTION_DEBOUNCE = 2000UL;  // Increased to 2 seconds for better reliability

// Function Prototypes
void initDisplay();
void showWelcome();
void showMainMenu();
void showEnterPassword();
void showUnlocked();
void showWrong();
void showLockedOut();
void handleKeypad();
void handleDoorLock();
void handlePeopleCounter();
void handleLighting();
void handleSmoke();
void ringSetColor(uint32_t c);
void ringClear();
void playTone(int freq, int ms);
void playStartupMelody();
void wheel(byte pos, uint8_t &r, uint8_t &g, uint8_t &b);

void setup() {
  Serial.begin(115200);
  Serial.println(F("=== Smart-Room Starting ==="));
  delay(100);

  Wire.begin();
  Serial.println(F("I2C initialized"));

  initDisplay();

  doorRing.begin();
  doorRing.setBrightness(80);
  ringClear();
  Serial.println(F("NeoPixel initialized"));

  // Pin modes
  pinMode(SOLENOID_RELAY_PIN, OUTPUT);
  pinMode(LIGHT_RELAY_PIN, OUTPUT);
  pinMode(BUZZER_PIN, OUTPUT);
  pinMode(OUTSIDE_IR, INPUT);
  pinMode(INSIDE_IR, INPUT);
  pinMode(SMOKE_DIGITAL, INPUT);
  pinMode(SMOKE_SW, INPUT_PULLUP);
  
  // Initial states
  digitalWrite(SOLENOID_RELAY_PIN, LOW);
  digitalWrite(LIGHT_RELAY_PIN, LOW);
  noTone(BUZZER_PIN);
  Serial.println(F("Pins initialized"));

  showWelcome();
  playStartupMelody();
  
  // Rainbow startup animation
  for (int j = 0; j < 256; j += 10) {
    for (int i = 0; i < DOOR_LEDS; i++) {
      uint8_t r, g, b;
      wheel((i * 256 / DOOR_LEDS + j) & 255, r, g, b);
      doorRing.setPixelColor(i, doorRing.Color(r, g, b));
    }
    doorRing.show();
    delay(10);
  }
  ringClear();

  Serial.println(F("System ready"));
  Serial.println(F(""));
  Serial.println(F("╔════════════════════════════════════╗"));
  Serial.println(F("║  Pin Configuration Summary:        ║"));
  Serial.println(F("║  Keypad Rows: 13, A2, A3, A4       ║"));
  Serial.println(F("║  Keypad Cols: A5, 10, 8, 7         ║"));
  Serial.println(F("║  Smoke Digital: 11                 ║"));
  Serial.println(F("║  Smoke Switch: 12                  ║"));
  Serial.println(F("╚════════════════════════════════════╝"));
  Serial.println(F(""));
}

void loop() {
  handleSmoke();          // Check smoke FIRST (highest priority)
  handleKeypad();
  handlePeopleCounter();
  handleDoorLock();
  handleLighting();

  display.clearDisplay();
  
  // Show fire warning if smoke detected
  if (smokeDetected) {
    display.setTextSize(2);
    display.setCursor(10, 5);
    display.println(F("!FIRE!"));
    display.setTextSize(1);
    display.setCursor(0, 25);
    display.println(F("Door: AUTO-OPEN"));
    display.println(F("Lights: EMERGENCY"));
    display.print(F("People: "));
    display.println(peopleCount);
  } else {
    switch (doorState) {
      case MAIN_MENU: showMainMenu(); break;
      case ENTER_PASSWORD: showEnterPassword(); break;
      case UNLOCKED: showUnlocked(); break;
      case WRONG: showWrong(); break;
      case LOCKED_OUT: showLockedOut(); break;
    }
  }
  
  display.display();
  delay(50);
}

// Display Helpers
void initDisplay() {
  Serial.println(F("Attempting OLED init..."));
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 init failed"));
    return;
  }
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  Serial.println(F("SSD1306 initialized"));
}

void showWelcome() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(15, 10);
  display.println(F("SMART"));
  display.setCursor(25, 30);
  display.println(F("ROOM"));
  display.setTextSize(1);
  display.setCursor(20, 50);
  display.println(F("Press A to begin"));
  display.display();
  delay(2000);
}

void showMainMenu() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.print(F("Door: "));
  display.println(doorUnlocked ? F("OPEN") : F("LOCKED"));
  
  display.setTextSize(2);
  display.setCursor(0, 15);
  display.print(F("People: "));
  display.println(peopleCount);
  
  display.setTextSize(1);
  display.setCursor(0, 35);
  display.print(F("Lights: "));
  display.println(lightsOn ? F("ON") : F("OFF"));
  
  // Show lux estimate using proper mapping
  int ldrRaw = analogRead(LIGHT_SENSOR);
  long lux = map(ldrRaw, 0, 1023, 0, 100000);
  display.print(F("Light: "));
  if (lux < 1000) {
    display.print(lux);
    display.println(F(" lux"));
  } else {
    display.print(lux / 1000);
    display.println(F("k lux"));
  }
  
  display.setCursor(0, 55);
  display.println(F("A=Pass B=+1 D=-1"));
}

void showEnterPassword() {
  display.setTextSize(1);
  display.setCursor(0, 0);
  display.println(F("Enter 4-digit code:"));
  display.setTextSize(2);
  char masked[5] = {0};
  for (size_t i = 0; i < strlen(enteredPassword); i++) masked[i] = '*';
  display.setCursor(20, 20);
  display.println(masked);
  display.setTextSize(1);
  display.setCursor(0, 50);
  display.println(F("#=OK D=Del C=Cancel"));
}

void showUnlocked() {
  display.setTextSize(2);
  display.setCursor(10, 10);
  display.println(F("UNLOCKED"));
  display.setTextSize(1);
  int secs = (UNLOCK_DURATION - (millis() - unlockTime)) / 1000;
  display.setCursor(10, 35);
  display.print(F("Auto-lock: "));
  display.print(secs);
  display.println(F("s"));
  display.setCursor(0, 50);
  display.println(F("C=Lock Now"));
}

void showWrong() {
  display.setTextSize(2);
  display.setCursor(20, 10);
  display.println(F("WRONG!"));
  display.setTextSize(1);
  display.setCursor(0, 35);
  display.print(F("Attempts: "));
  display.print(wrongAttempts);
  display.print(F("/"));
  display.println(MAX_ATTEMPTS);
}

void showLockedOut() {
  display.setTextSize(2);
  display.setCursor(5, 10);
  display.println(F("LOCKED"));
  display.println(F("  OUT!"));
  display.setTextSize(1);
  unsigned long remaining = (lockoutStartTime + LOCKOUT_DURATION - millis()) / 1000;
  display.setCursor(20, 45);
  display.print(F("Wait "));
  display.print(remaining);
  display.println(F("s"));
}

// Input Handling
void handleKeypad() {
  char key = keypad.getKey();
  if (!key) return;

  Serial.print(F("Key: "));
  Serial.println(key);

  // Fire test button (always available)
  if (key == '*') {
    smokeDetected = !smokeDetected;
    Serial.print(F("Manual fire test: "));
    Serial.println(smokeDetected ? "ON" : "OFF");
    if (!smokeDetected) {
      noTone(BUZZER_PIN);
      ringClear();
      delay(50);
    }
    return;
  }

  // Manual people counter buttons (only in main menu)
  if (key == 'B' && doorState == MAIN_MENU) {
    peopleCount++;
    Serial.print(F("Manual +1, Total: "));
    Serial.println(peopleCount);
    playTone(1200, 100);
    return;
  }
  
  if (key == 'D' && doorState == MAIN_MENU && peopleCount > 0) {
    peopleCount--;
    Serial.print(F("Manual -1, Total: "));
    Serial.println(peopleCount);
    playTone(800, 100);
    return;
  }

  // Start password entry
  if (doorState == MAIN_MENU && key == 'A') {
    doorState = ENTER_PASSWORD;
    enteredPassword[0] = '\0';
    return;
  }

  // Password entry mode
  if (doorState == ENTER_PASSWORD) {
    if (key >= '0' && key <= '9' && strlen(enteredPassword) < 4) {
      size_t len = strlen(enteredPassword);
      enteredPassword[len] = key;
      enteredPassword[len + 1] = '\0';
      playTone(1200, 80);
    }
    else if (key == 'D' && strlen(enteredPassword) > 0) {
      enteredPassword[strlen(enteredPassword) - 1] = '\0';
      playTone(800, 80);
    }
    else if (key == '#' && strlen(enteredPassword) == 4) {
      if (strcmp(enteredPassword, CORRECT_PASSWORD) == 0) {
        digitalWrite(SOLENOID_RELAY_PIN, HIGH);
        doorUnlocked = true;
        unlockTime = millis();
        wrongAttempts = 0;
        doorState = UNLOCKED;
        ringSetColor(0x00FF00);
        playTone(1500, 200);
        delay(100);
        playTone(1800, 200);
        Serial.println(F("✓ Door Unlocked"));
      } else {
        wrongAttempts++;
        doorState = WRONG;
        ringSetColor(0xFF0000);
        playTone(400, 500);
        Serial.println(F("✗ Wrong password"));
        delay(1500);
        
        if (wrongAttempts >= MAX_ATTEMPTS) {
          doorState = LOCKED_OUT;
          lockoutStartTime = millis();
          ringClear();
          for (int i = 0; i < 6; i++) {
            ringSetColor(0xFF0000);
            delay(150);
            ringClear();
            delay(150);
          }
          Serial.println(F("⚠ Locked out"));
        } else {
          doorState = MAIN_MENU;
          ringClear();
        }
      }
      enteredPassword[0] = '\0';
    }
    else if (key == 'C') {
      doorState = MAIN_MENU;
      enteredPassword[0] = '\0';
      ringClear();
    }
    return;
  }

  // Manual lock from unlocked state
  if (doorState == UNLOCKED && key == 'C') {
    digitalWrite(SOLENOID_RELAY_PIN, LOW);
    doorUnlocked = false;
    doorState = MAIN_MENU;
    ringClear();
    Serial.println(F("🔒 Door manually locked"));
    return;
  }
}

void handleDoorLock() {
  // Don't auto-lock during fire emergency
  if (smokeDetected) return;
  
  // Auto-lock after timeout
  if (doorState == UNLOCKED && (millis() - unlockTime >= UNLOCK_DURATION)) {
    digitalWrite(SOLENOID_RELAY_PIN, LOW);
    doorUnlocked = false;
    doorState = MAIN_MENU;
    ringClear();
    playTone(800, 150);
    Serial.println(F("🔒 Door auto-locked"));
  }

  // End lockout period
  if (doorState == LOCKED_OUT && (millis() - lockoutStartTime >= LOCKOUT_DURATION)) {
    doorState = MAIN_MENU;
    wrongAttempts = 0;
    Serial.println(F("✓ Lockout period ended"));
  }
}

void handlePeopleCounter() {
  bool outsidePIR = digitalRead(OUTSIDE_IR);
  bool insidePIR = digitalRead(INSIDE_IR);
  unsigned long now = millis();

  // Debug output every 5 seconds
  static unsigned long lastPrint = 0;
  if (now - lastPrint >= 5000) {
    Serial.print(F("👥 PIR - Outside: "));
    Serial.print(outsidePIR ? "HIGH" : "LOW");
    Serial.print(F(", Inside: "));
    Serial.print(insidePIR ? "HIGH" : "LOW");
    Serial.print(F(", People Count: "));
    Serial.println(peopleCount);
    lastPrint = now;
  }

  // OUTSIDE PIR: Person entering (rising edge detection)
  if (outsidePIR && !lastOutsidePIR && (now - lastOutsideMotionTime >= MOTION_DEBOUNCE)) {
    peopleCount++;
    lastOutsideMotionTime = now;
    playTone(1200, 100);
    Serial.print(F("✓ Outside motion! Person ENTERED. Total: "));
    Serial.println(peopleCount);
  }

  // INSIDE PIR: Person exiting (rising edge detection)
  if (insidePIR && !lastInsidePIR && (now - lastInsideMotionTime >= MOTION_DEBOUNCE) && peopleCount > 0) {
    peopleCount--;
    lastInsideMotionTime = now;
    playTone(900, 100);
    Serial.print(F("✓ Inside motion! Person EXITED. Total: "));
    Serial.println(peopleCount);
  }

  lastOutsidePIR = outsidePIR;
  lastInsidePIR = insidePIR;
}

void handleLighting() {
  // Emergency lighting during fire
  if (smokeDetected) {
    if (!lightsOn) {
      digitalWrite(LIGHT_RELAY_PIN, HIGH);
      lightsOn = true;
      Serial.println(F("💡 Lights: EMERGENCY ON"));
    }
    return;
  }

  int ldrRaw = analogRead(LIGHT_SENSOR);
  
  // Map analog value (0-1023) to lux (0-100000)
  // In Wokwi LDR: Low analog = Dark, High analog = Bright
  long estimatedLux = map(ldrRaw, 0, 1023, 0, 100000);
  
  // Use constant for threshold
  bool isDark = (estimatedLux < DARK_THRESHOLD_LUX);
  bool isOccupied = (peopleCount > 0);

  // Debug output every 2 seconds
  static unsigned long lastLightPrint = 0;
  if (millis() - lastLightPrint >= 2000) {
    Serial.print(F("💡 LDR: "));
    Serial.print(ldrRaw);
    Serial.print(F(" → ~"));
    Serial.print(estimatedLux);
    Serial.print(F(" lux"));
    Serial.print(isDark ? F(" (DARK)") : F(" (BRIGHT)"));
    Serial.print(F(", People: "));
    Serial.print(peopleCount);
    Serial.print(F(", Lights: "));
    Serial.println(lightsOn ? "ON" : "OFF");
    lastLightPrint = millis();
  }

  // Turn lights ON if dark AND occupied
  if (isDark && isOccupied && !lightsOn) {
    digitalWrite(LIGHT_RELAY_PIN, HIGH);
    lightsOn = true;
    Serial.println(F("💡 Lights: AUTO ON (dark + occupied)"));
  }
  // Turn lights OFF if not dark OR not occupied
  else if ((!isDark || !isOccupied) && lightsOn) {
    digitalWrite(LIGHT_RELAY_PIN, LOW);
    lightsOn = false;
    Serial.print(F("💡 Lights: AUTO OFF ("));
    if (!isDark) Serial.print(F("bright"));
    if (!isDark && !isOccupied) Serial.print(F(" + "));
    if (!isOccupied) Serial.print(F("empty"));
    Serial.println(F(")"));
  }
}

void handleSmoke() {
  int smokeAnalog = analogRead(SMOKE_ANALOG);
  bool smokeDigital = digitalRead(SMOKE_DIGITAL);
  bool smokeSwitch = !digitalRead(SMOKE_SW);  // Inverted (pullup)
  
  // Fire detected if ANY condition is true
  bool fireDetected = (smokeAnalog >= SMOKE_THRESHOLD) || smokeDigital || smokeSwitch;

  // Debug output every 1 second
  static unsigned long lastSmokePrint = 0;
  if (millis() - lastSmokePrint >= 1000) {
    Serial.print(F("🔥 Smoke Sensor - Analog: "));
    Serial.print(smokeAnalog);
    Serial.print(F(" (threshold: "));
    Serial.print(SMOKE_THRESHOLD);
    Serial.print(F("), Digital: "));
    Serial.print(smokeDigital);
    Serial.print(F(", Switch: "));
    Serial.print(smokeSwitch);
    Serial.print(F(" → Fire: "));
    Serial.println(fireDetected ? "⚠️ DETECTED" : "✓ Clear");
    lastSmokePrint = millis();
  }

  // FIRE STARTED
  if (fireDetected && !smokeDetected) {
    smokeDetected = true;
    Serial.println(F(""));
    Serial.println(F("╔═══════════════════════════════╗"));
    Serial.println(F("║   🔥🔥 FIRE DETECTED! 🔥🔥    ║"));
    Serial.println(F("╚═══════════════════════════════╝"));
    
    // Emergency unlock door
    digitalWrite(SOLENOID_RELAY_PIN, HIGH);
    doorUnlocked = true;
    unlockTime = millis();
    doorState = UNLOCKED;
    
    // Turn on emergency lights
    digitalWrite(LIGHT_RELAY_PIN, HIGH);
    lightsOn = true;
    
    // Visual alarm
    ringSetColor(0xFF0000);
    
    Serial.println(F("→ Door: AUTO-UNLOCKED ✓"));
    Serial.println(F("→ Lights: EMERGENCY ON ✓"));
    Serial.println(F("→ Alarm: SOUNDING ✓"));
    Serial.println(F(""));
  }
  // FIRE CLEARED
  else if (!fireDetected && smokeDetected) {
    smokeDetected = false;
    
    // Immediately stop all alarms
    noTone(BUZZER_PIN);
    ringClear();
    
    Serial.println(F(""));
    Serial.println(F("╔═══════════════════════════════╗"));
    Serial.println(F("║   ✓✓ FIRE CLEARED! ✓✓        ║"));
    Serial.println(F("╚═══════════════════════════════╝"));
    Serial.println(F("→ Buzzer: OFF ✓"));
    Serial.println(F("→ LED Ring: CLEARED ✓"));
    Serial.println(F("→ Returning to normal mode"));
    Serial.println(F(""));
    
    delay(100);  // Ensure commands are processed
  }

  // Fire alarm sound: alternating high/low tones
  if (smokeDetected) {
    if (millis() - fireAlarmTimer >= 250) {
      fireAlarmHigh = !fireAlarmHigh;
      tone(BUZZER_PIN, fireAlarmHigh ? 1000 : 600);
      fireAlarmTimer = millis();
      
      // Flash red ring
      ringSetColor(fireAlarmHigh ? 0xFF0000 : 0x550000);
    }
  } else {
    // Ensure alarm is OFF when no smoke
    noTone(BUZZER_PIN);
    if (!doorUnlocked) {
      ringClear();
    }
  }
}

void ringSetColor(uint32_t c) {
  for (int i = 0; i < DOOR_LEDS; i++) doorRing.setPixelColor(i, c);
  doorRing.show();
}

void ringClear() {
  ringSetColor(0);
}

void playTone(int freq, int ms) {
  tone(BUZZER_PIN, freq, ms);
  delay(ms + 10);
  noTone(BUZZER_PIN);
}

void playStartupMelody() {
  int notes[] = {262, 330, 392, 523};
  for (int i = 0; i < 4; i++) {
    playTone(notes[i], 150);
  }
}

void wheel(byte pos, uint8_t &r, uint8_t &g, uint8_t &b) {
  pos = 255 - pos;
  if (pos < 85) { 
    r = 255 - pos * 3; 
    g = 0; 
    b = pos * 3; 
  }
  else if (pos < 170) { 
    pos -= 85; 
    r = 0; 
    g = pos * 3; 
    b = 255 - pos * 3; 
  }
  else { 
    pos -= 170; 
    r = pos * 3; 
    g = 255 - pos * 3; 
    b = 0; 
  }
}