#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>

// OLED display dimensions
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// I2C pins for ESP32-C3
#define OLED_SDA 4
#define OLED_SCL 5
#define OLED_ADDR 0x3C

// Button and LED pins
#define BUTTON_PIN 2
#define LED_PIN 10

// Create display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Screen states for automatic demonstration
enum ScreenState {
  OFF_STATE,           // Screen OFF - waiting for button press
  WAKE_SCREEN,        // Screen 0: "Hello! I'm E.L.I."
  HOME_MENU,          // Screen 1: Home with 3 options
  START_FOCUS_MENU,   // Screen 1A: Start Focus submenu
  POMODORO_TIMER,     // Screen 1A-i: Pomodoro countdown
  CUSTOM_TIMER,       // Custom timer screen
  FOCUS_BLUEPRINT,    // Focus Blueprint screen
  CREATE_TIMER_MENU,  // Create Timer screen
  END_SCREEN          // Final screen before turning off
};

ScreenState currentState = OFF_STATE;
unsigned long stateStartTime = 0;
int demoCycle = 0;

// For countdown simulation
int countdownSeconds = 25;
int countdownMinutes = 0;
unsigned long lastCountdownUpdate = 0;

bool demoActive = false;

void setup() {
  Serial.begin(115200);
  
  pinMode(BUTTON_PIN, INPUT_PULLUP);
  pinMode(LED_PIN, OUTPUT);
  digitalWrite(LED_PIN, LOW);  // LED starts OFF
  
  Wire.begin(OLED_SDA, OLED_SCL);
  if(!display.begin(SSD1306_SWITCHCAPVCC, OLED_ADDR)) {
    Serial.println("OLED initialization failed!");
  }
  
  // Clear display buffer completely
  display.clearDisplay();
  display.display();
  
  Serial.println("=== UI DEMO READY ===");
  Serial.println("Press button to start");
}

void loop() {
  if (digitalRead(BUTTON_PIN) == LOW) {
    delay(50);
    if (digitalRead(BUTTON_PIN) == LOW) {
      if (!demoActive) {
        startDemo();
      } else {
        restartDemo();
      }
      while(digitalRead(BUTTON_PIN) == LOW) {
        delay(10);
      }
    }
  }
  
  if (demoActive) {
    if (millis() - stateStartTime >= 2000) {
      // After END_SCREEN, turn off instead of moving to next screen
      if (currentState == END_SCREEN) {
        endDemo();
      } else {
        moveToNextScreen();
        stateStartTime = millis();
      }
    }
    displayCurrentScreen();
  }
  
  delay(50);
}

void startDemo() {
  demoActive = true;
  currentState = WAKE_SCREEN;
  stateStartTime = millis();
  demoCycle = 0;
  countdownSeconds = 25;
  countdownMinutes = 0;
  
  // 🔴 TURN LED ON when button is pressed
  digitalWrite(LED_PIN, HIGH);
  
  Serial.println("=== DEMO STARTED ===");
}

void moveToNextScreen() {
  switch(currentState) {
    case WAKE_SCREEN:
      currentState = HOME_MENU;
      Serial.println("→ HOME MENU");
      break;
      
    case HOME_MENU:
      if (demoCycle == 0) {
        currentState = START_FOCUS_MENU;
        Serial.println("→ START FOCUS");
      } else if (demoCycle == 1) {
        currentState = CREATE_TIMER_MENU;
        Serial.println("→ CREATE TIMER");
      } else {
        currentState = FOCUS_BLUEPRINT;
        Serial.println("→ FOCUS BLUEPRINT");
      }
      demoCycle = (demoCycle + 1) % 3;
      break;
      
    case START_FOCUS_MENU:
      if (demoCycle == 0) {
        currentState = POMODORO_TIMER;
        countdownSeconds = 25;
        countdownMinutes = 0;
        Serial.println("→ POMODORO");
      } else {
        currentState = CUSTOM_TIMER;
        Serial.println("→ CUSTOM TIMER");
      }
      break;
      
    case POMODORO_TIMER:
    case CUSTOM_TIMER:
    case FOCUS_BLUEPRINT:
    case CREATE_TIMER_MENU:
      currentState = END_SCREEN;
      Serial.println("→ END SCREEN");
      break;
      
    default:
      break;
  }
  
  // REMOVED the LED blink here - LED stays solid ON
}

void displayCurrentScreen() {
  display.clearDisplay();
  
  switch(currentState) {
    case WAKE_SCREEN:
      drawWakeScreen();
      break;
    case HOME_MENU:
      drawHomeMenu();
      break;
    case START_FOCUS_MENU:
      drawStartFocusMenu();
      break;
    case POMODORO_TIMER:
      drawPomodoroTimer();
      break;
    case CUSTOM_TIMER:
      drawCustomTimer();
      break;
    case FOCUS_BLUEPRINT:
      drawFocusBlueprint();
      break;
    case CREATE_TIMER_MENU:
      drawCreateTimerMenu();
      break;
    case END_SCREEN:
      drawEndScreen();
      break;
    default:
      return;
  }
  
  display.display();
}

// Screen 0: Wake Screen
void drawWakeScreen() {
  display.drawLine(20, 12, 108, 12, SSD1306_WHITE);
  display.drawLine(20, 52, 108, 52, SSD1306_WHITE);
  
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(32, 22);
  display.print("Hello");
  
  display.setTextSize(1);
  display.setCursor(40, 42);
  display.print("I'm E.L.I.");
}

// Screen 1: Home Menu
void drawHomeMenu() {
  display.setTextSize(1);
  display.setCursor(44, 4);
  display.print("HOME");
  display.drawLine(40, 14, 88, 14, SSD1306_WHITE);
  
  // Simple menu without arrows
  display.setCursor(25, 24);
  display.print("Start Focus");
  
  display.setCursor(25, 38);
  display.print("Create Timer");
  
  display.setCursor(25, 52);
  display.print("Focus Blueprint");
  
  // Simple navigation hint
  display.drawLine(0, 63, 127, 63, SSD1306_WHITE);
}

// Screen 1A: Start Focus Submenu
void drawStartFocusMenu() {
  display.setTextSize(1);
  display.setCursor(30, 4);
  display.print("START FOCUS");
  display.drawLine(25, 14, 103, 14, SSD1306_WHITE);
  
  display.setCursor(25, 26);
  display.print("Pomodoro (25+5)");
  
  display.setCursor(25, 40);
  display.print("Custom Timer");
  
  display.setCursor(25, 54);
  display.print("Back");
  
  display.drawLine(0, 63, 127, 63, SSD1306_WHITE);
}

// Pomodoro Timer
void drawPomodoroTimer() {
  if (millis() - lastCountdownUpdate > 500) {
    lastCountdownUpdate = millis();
    countdownSeconds--;
    if (countdownSeconds < 0) {
      countdownSeconds = 25;
      // REMOVED the LED flashing here - LED stays solid ON
    }
  }
  
  display.setTextSize(1);
  display.setCursor(32, 4);
  display.print("FOCUS TIME");
  display.drawLine(25, 14, 103, 14, SSD1306_WHITE);
  
  display.setTextSize(3);
  display.setCursor(32, 22);
  
  if (countdownSeconds < 10) {
    display.print(countdownMinutes);
    display.print(":0");
    display.print(countdownSeconds);
  } else {
    display.print(countdownMinutes);
    display.print(":");
    display.print(countdownSeconds);
  }
  
  display.drawLine(0, 50, 127, 50, SSD1306_WHITE);
  display.setTextSize(1);
  display.setCursor(30, 54);
  display.print("Stand & stretch!");
  
  if (countdownSeconds == 0 && countdownMinutes == 0) {
    display.fillRect(40, 56, 48, 8, SSD1306_WHITE);
    display.setTextColor(SSD1306_BLACK);
    display.setCursor(45, 56);
    display.print("DONE");
    display.setTextColor(SSD1306_WHITE);
  }
}

// Custom Timer
void drawCustomTimer() {
  display.setTextSize(1);
  display.setCursor(35, 4);
  display.print("CUSTOM");
  display.drawLine(30, 14, 98, 14, SSD1306_WHITE);
  
  display.setCursor(30, 26);
  display.print("5 minutes");
  
  display.setCursor(30, 40);
  display.print("10 minutes");
  
  display.setCursor(30, 54);
  display.print("15 minutes");
  
  display.drawLine(0, 63, 127, 63, SSD1306_WHITE);
}

// Focus Blueprint
void drawFocusBlueprint() {
  display.setTextSize(1);
  display.setCursor(20, 4);
  display.print("FOCUS BLUEPRINT");
  display.drawLine(15, 14, 113, 14, SSD1306_WHITE);
  
  display.setCursor(20, 26);
  display.print("Focus: 85%");
  
  display.setCursor(20, 40);
  display.print("Breaks: 4");
  
  display.setCursor(20, 54);
  display.print("Productivity: High");
  
  display.drawLine(0, 63, 127, 63, SSD1306_WHITE);
}

// Create Timer Menu
void drawCreateTimerMenu() {
  display.setTextSize(1);
  display.setCursor(25, 4);
  display.print("CREATE TIMER");
  display.drawLine(20, 14, 108, 14, SSD1306_WHITE);
  
  display.setCursor(30, 26);
  display.print("Work: 25m");
  
  display.setCursor(30, 40);
  display.print("Break: 5m");
  
  display.drawLine(20, 50, 108, 50, SSD1306_WHITE);
  display.setCursor(45, 54);
  display.print("SAVE");
  
  display.drawLine(0, 63, 127, 63, SSD1306_WHITE);
}

// End Screen - Shows before turning off
void drawEndScreen() {
  display.setTextSize(1);
  display.setCursor(25, 22);
  display.print("Demo Complete");
  
  display.setCursor(35, 38);
  display.print("Turning Off");
  
  // Simple dots animation
  static int dotCount = 0;
  dotCount = (dotCount + 1) % 4;
  display.setCursor(55, 52);
  for(int i = 0; i < dotCount; i++) {
    display.print(".");
  }
}

void endDemo() {
  demoActive = false;
  currentState = OFF_STATE;
  
  display.clearDisplay();
  display.display();
  
  // 🔴 TURN LED OFF when OLED turns off
  digitalWrite(LED_PIN, LOW);
  
  Serial.println("=== DEMO ENDED ===");
  Serial.println("Press button to start again");
}

void restartDemo() {
  display.clearDisplay();
  display.display();
  delay(200);
  
  demoActive = true;
  currentState = WAKE_SCREEN;
  stateStartTime = millis();
  demoCycle = 0;
  countdownSeconds = 25;
  countdownMinutes = 0;
  
  // LED is already ON from startDemo, so no change here
  
  Serial.println("=== DEMO RESTARTED ===");
}