#include <Wire.h>
#include <U8g2lib.h>

// Pin definitions for Metro
#define LED_PIN 9
#define BUZZER_PIN 3
#define FLEX_PIN A2

// U8g2 object for 128x64 OLED with I2C
U8G2_SSD1306_128X64_NONAME_1_HW_I2C display(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);

// Flex calibration (adjust after testing)
const int flatValue = 621;
const int bentValue = 660;
const int totalRange = bentValue - flatValue;
const int scrollThreshold = flatValue + (totalRange * 0.3);
const int selectThreshold = flatValue + (totalRange * 0.7);

// Menu states
enum ScreenState {
  WAKE_SCREEN,
  HOME_MENU,
  START_FOCUS_MENU,
  POMODORO_TIMER,
  CUSTOM_TIMER_PRESETS,
  FOCUS_BLUEPRINT,
  EDIT_CUSTOM_TIMER
};

ScreenState currentState = WAKE_SCREEN;
int menuCursor = 0;
unsigned long lastInteraction = 0;

// Gesture variables
int lastBendLevel = 0;
unsigned long lastGestureTime = 0;
unsigned long pressStartTime = 0;
bool isLongPress = false;
bool isVeryLongPress = false;          // new flag for very long press (power off)
const unsigned long longPressDuration = 1000;      // 1 sec for select
const unsigned long veryLongPressDuration = 3000;  // 3 sec for power off
const unsigned long gestureDebounce = 300;          // 300 ms

// Double‑press window
const unsigned long doublePressWindow = 600;

// Custom timer
int customHours = 0;
int customMinutes = 25;
int editMode = 0;
const int presetTimes[] = {5, 10, 15};
const int numPresets = 3;

// Function prototypes
void drawWakeScreen();
void drawHomeMenu();
void drawStartFocusMenu();
void drawPomodoroTimer();
void drawCustomTimerPresets();
void drawFocusBlueprint();
void drawEditCustomTimer();
void handleScroll();
void handleSelect();
void handleBack();
void handlePowerOff();                   // new function
void handleGestures(int currentBendLevel);
void updateDisplay();

void setup() {
  Serial.begin(9600);
  while (!Serial); // optional
  Serial.println("=== Metro Flex Menu (U8g2) ===");

  pinMode(LED_PIN, OUTPUT);
  pinMode(BUZZER_PIN, OUTPUT);

  Wire.begin();
  delay(500);

  display.begin();
  display.clear();
  display.setFont(u8g2_font_ncenB08_tr);

  digitalWrite(LED_PIN, HIGH);
  tone(BUZZER_PIN, 1000, 200);
  delay(300);
  digitalWrite(LED_PIN, LOW);

  drawWakeScreen();
  Serial.println("Ready. Bend sensor to navigate.");
}

void loop() {
  int flexValue = analogRead(FLEX_PIN);
  int percent = map(flexValue, flatValue, bentValue, 0, 100);
  percent = constrain(percent, 0, 100);

  int bendLevel = 0;
  if (flexValue >= selectThreshold) bendLevel = 2;
  else if (flexValue >= scrollThreshold) bendLevel = 1;

  handleGestures(bendLevel);
  updateDisplay();

  if (millis() - lastInteraction > 60000) {
    currentState = WAKE_SCREEN;
    menuCursor = 0;
  }

  delay(100);
}

void handleGestures(int currentBendLevel) {
  static unsigned long pressTimer = 0;

  if (currentBendLevel > 0 && lastBendLevel == 0) {
    pressTimer = millis();
    isLongPress = false;
    isVeryLongPress = false;
    Serial.println("Bend start");
  }

  // Check for very long press (power off) – takes priority over long press
  if (currentBendLevel > 0 && !isVeryLongPress && (millis() - pressTimer > veryLongPressDuration)) {
    isVeryLongPress = true;
    if (millis() - lastGestureTime > gestureDebounce) {
      lastGestureTime = millis();
      Serial.println("VERY LONG PRESS (POWER OFF)");
      handlePowerOff();
    }
  }
  // Check for long press (select) – only if very long not triggered
  else if (currentBendLevel > 0 && !isLongPress && !isVeryLongPress && (millis() - pressTimer > longPressDuration)) {
    isLongPress = true;
    if (millis() - lastGestureTime > gestureDebounce) {
      lastGestureTime = millis();
      Serial.println("LONG PRESS (SELECT)");
      digitalWrite(LED_PIN, HIGH);
      tone(BUZZER_PIN, 1500, 200);
      delay(200);
      digitalWrite(LED_PIN, LOW);
      handleSelect();
    }
  }

  // Release
  if (currentBendLevel == 0 && lastBendLevel > 0) {
    if (!isLongPress && !isVeryLongPress && (millis() - pressTimer) < longPressDuration) {
      if (millis() - lastGestureTime > gestureDebounce) {
        lastGestureTime = millis();
        Serial.println("SHORT PRESS (SCROLL)");
        digitalWrite(LED_PIN, HIGH);
        tone(BUZZER_PIN, 800, 50);
        delay(20);
        digitalWrite(LED_PIN, LOW);
        handleScroll();
      }
    }
  }

  // Double press detection (only if not in a long press)
  static unsigned long lastReleaseTime = 0;
  static int releaseCount = 0;

  if (currentBendLevel == 0 && lastBendLevel > 0) {
    if (millis() - lastReleaseTime < doublePressWindow) {
      releaseCount++;
      if (releaseCount >= 2) {
        if (millis() - lastGestureTime > gestureDebounce) {
          lastGestureTime = millis();
          Serial.println("DOUBLE PRESS (BACK)");
          digitalWrite(LED_PIN, HIGH);
          tone(BUZZER_PIN, 600, 50);
          delay(100);
          tone(BUZZER_PIN, 600, 50);
          delay(50);
          digitalWrite(LED_PIN, LOW);
          handleBack();
          releaseCount = 0;
        }
      }
    } else {
      releaseCount = 1;
    }
    lastReleaseTime = millis();
  }

  lastBendLevel = currentBendLevel;
}

void handlePowerOff() {
  // Visual/audio feedback for power off
  digitalWrite(LED_PIN, HIGH);
  tone(BUZZER_PIN, 400, 800);  // low long beep
  delay(800);
  digitalWrite(LED_PIN, LOW);

  // Return to wake screen and reset cursor
  currentState = WAKE_SCREEN;
  menuCursor = 0;
  lastInteraction = millis();  // reset idle timer

  Serial.println("Power off – returning to wake screen");
}

void handleScroll() {
  lastInteraction = millis();
  int maxItems;
  switch (currentState) {
    case HOME_MENU:
      maxItems = 3;
      menuCursor = (menuCursor + 1) % maxItems;
      break;
    case START_FOCUS_MENU:
      maxItems = 3; // Pomodoro, Custom Timer, Back
      menuCursor = (menuCursor + 1) % maxItems;
      break;
    case CUSTOM_TIMER_PRESETS:
      maxItems = numPresets + 1; // presets + Back
      menuCursor = (menuCursor + 1) % maxItems;
      break;
    case FOCUS_BLUEPRINT:
      maxItems = 2;
      menuCursor = (menuCursor + 1) % maxItems;
      break;
    case POMODORO_TIMER:
      menuCursor = (menuCursor + 1) % 2;
      break;
    case EDIT_CUSTOM_TIMER:
      menuCursor = (menuCursor + 1) % 2;
      break;
    default:
      break;
  }
}

void handleSelect() {
  lastInteraction = millis();
  switch (currentState) {
    case WAKE_SCREEN:
      currentState = HOME_MENU;
      menuCursor = 0;
      break;

    case HOME_MENU:
      if (menuCursor == 0) currentState = START_FOCUS_MENU;
      else if (menuCursor == 1) currentState = CUSTOM_TIMER_PRESETS;
      else if (menuCursor == 2) currentState = FOCUS_BLUEPRINT;
      menuCursor = 0;
      break;

    case START_FOCUS_MENU:
      if (menuCursor == 0) currentState = POMODORO_TIMER;
      else if (menuCursor == 1) {
        currentState = EDIT_CUSTOM_TIMER;
        customHours = 0;
        customMinutes = 25;
        editMode = 0;
      } else if (menuCursor == 2) { // Back option
        handleBack();
      }
      menuCursor = 0;
      break;

    case CUSTOM_TIMER_PRESETS:
      if (menuCursor < numPresets) {
        currentState = EDIT_CUSTOM_TIMER;
        customHours = 0;
        customMinutes = presetTimes[menuCursor];
        editMode = 0;
      } else { // last item = Back
        handleBack();
      }
      menuCursor = 0;
      break;

    case FOCUS_BLUEPRINT:
      if (menuCursor == 0) { // Back
        handleBack();
      }
      break;

    case POMODORO_TIMER:
      if (menuCursor == 0) { // Back
        handleBack();
      }
      break;

    case EDIT_CUSTOM_TIMER:
      if (menuCursor == 0) {
        // Save and start timer
        currentState = POMODORO_TIMER;
      } else { // Back
        handleBack();
      }
      menuCursor = 0;
      break;

    default: break;
  }
}

void handleBack() {
  lastInteraction = millis();
  switch (currentState) {
    case HOME_MENU:
      currentState = WAKE_SCREEN;
      break;
    case START_FOCUS_MENU:
    case CUSTOM_TIMER_PRESETS:
    case FOCUS_BLUEPRINT:
      currentState = HOME_MENU;
      menuCursor = 1; // optional
      break;
    case POMODORO_TIMER:
      currentState = START_FOCUS_MENU;
      menuCursor = 0;
      break;
    case EDIT_CUSTOM_TIMER:
      currentState = START_FOCUS_MENU;
      menuCursor = 0;
      break;
    default: break;
  }
}

// U8g2 display functions
void updateDisplay() {
  display.firstPage();
  do {
    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_PRESETS: drawCustomTimerPresets(); break;
      case FOCUS_BLUEPRINT: drawFocusBlueprint(); break;
      case EDIT_CUSTOM_TIMER: drawEditCustomTimer(); break;
    }
    // Back arrow hint on all screens except wake
    if (currentState != WAKE_SCREEN) {
      display.setFont(u8g2_font_ncenB08_tr);
      display.drawStr(100, 60, "<<");
    }
  } while (display.nextPage());
}

void drawWakeScreen() {
  display.setFont(u8g2_font_ncenB12_tr);
  display.drawStr(32, 30, "Hello");
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(40, 50, "I'm E.L.I.");
}

void drawHomeMenu() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(44, 10, "HOME");
  display.drawLine(40, 14, 88, 14);

  const char* items[] = {"Start Focus", "Create Timer", "Focus Blueprint"};
  for (int i = 0; i < 3; i++) {
    int y = 24 + i * 14;
    if (i == menuCursor) display.drawStr(15, y, ">");
    display.drawStr(25, y, items[i]);
  }
}

void drawStartFocusMenu() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(30, 10, "START FOCUS");
  display.drawLine(25, 14, 103, 14);

  const char* items[] = {"Pomodoro (25+5)", "Custom Timer", "← Back"};
  for (int i = 0; i < 3; i++) {
    int y = 26 + i * 14;
    if (i == menuCursor) display.drawStr(15, y, ">");
    display.drawStr(25, y, items[i]);
  }
}

void drawPomodoroTimer() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(32, 10, "FOCUS TIME");
  display.setFont(u8g2_font_logisoso32_tf);
  display.drawStr(20, 45, "25:00");
  display.setFont(u8g2_font_ncenB08_tr);
  if (menuCursor == 0) {
    display.drawStr(15, 60, ">");
  }
  display.drawStr(25, 60, "← Back");
}

void drawCustomTimerPresets() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(30, 10, "CUSTOM TIMER");
  display.drawLine(20, 14, 108, 14);

  for (int i = 0; i < numPresets; i++) {
    int y = 30 + i * 14;
    if (i == menuCursor) display.drawStr(15, y, ">");
    char buf[10];
    sprintf(buf, "%d min", presetTimes[i]);
    display.drawStr(25, y, buf);
  }
  int yBack = 30 + numPresets * 14;
  if (menuCursor == numPresets) display.drawStr(15, yBack, ">");
  display.drawStr(25, yBack, "← Back");
}

void drawFocusBlueprint() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(20, 10, "Your Blueprint");
  display.drawLine(15, 14, 113, 14);
  display.drawStr(20, 30, "Sessions: 12");
  display.drawStr(20, 44, "Best: 6 AM");
  if (menuCursor == 0) display.drawStr(15, 60, ">");
  display.drawStr(25, 60, "← Back");
}

void drawEditCustomTimer() {
  display.setFont(u8g2_font_ncenB08_tr);
  display.drawStr(15, 10, "New Task Timer");
  display.drawLine(15, 14, 113, 14);
  display.setFont(u8g2_font_ncenB12_tr);

  char buf[20];
  if (editMode == 0)
    sprintf(buf, "[%d]:%02d", customHours, customMinutes);
  else
    sprintf(buf, "%d:[%02d]", customHours, customMinutes);
  display.drawStr(25, 35, buf);

  display.setFont(u8g2_font_ncenB08_tr);
  if (menuCursor == 0) display.drawStr(15, 50, ">");
  display.drawStr(25, 50, "Save & Start");
  if (menuCursor == 1) display.drawStr(15, 60, ">");
  display.drawStr(25, 60, "← Back");
}