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

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

// Initialize OLED display
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Button pins
#define UP_BUTTON_PIN1 6
#define UP_BUTTON_PIN2 17
#define DOWN_BUTTON_PIN1 0
#define DOWN_BUTTON_PIN2 2
#define LEFT_BUTTON_PIN1 16
#define LEFT_BUTTON_PIN2 4
#define RIGHT_BUTTON_PIN1 19
#define RIGHT_BUTTON_PIN2 18
#define SELECT_BUTTON_PIN1 1  // D1
#define SELECT_BUTTON_PIN2 0  // D0

// Menu states
enum MenuState {
  MAIN_MENU,
  SENSOR_DISPLAY,
  LED_CONTROL,
  SETTINGS
};

// Current menu state
MenuState currentState = MAIN_MENU;

// Variables for LED control
bool ledState = false;
int ledBrightness = 128;

// Variables for sensor simulation
float temperature = 25.0;
float humidity = 50.0;

// Menu navigation
int menuIndex = 0;
const char* menuItems[] = {"Sensor Display", "LED Control", "Settings", "Info"};

void setup() {
  Serial.begin(115200);
  
  // Initialize buttons
  pinMode(UP_BUTTON_PIN1, INPUT_PULLUP);
  pinMode(UP_BUTTON_PIN2, INPUT_PULLUP);
  pinMode(DOWN_BUTTON_PIN1, INPUT_PULLUP);
  pinMode(DOWN_BUTTON_PIN2, INPUT_PULLUP);
  pinMode(LEFT_BUTTON_PIN1, INPUT_PULLUP);
  pinMode(LEFT_BUTTON_PIN2, INPUT_PULLUP);
  pinMode(RIGHT_BUTTON_PIN1, INPUT_PULLUP);
  pinMode(RIGHT_BUTTON_PIN2, INPUT_PULLUP);
  pinMode(SELECT_BUTTON_PIN1, INPUT_PULLUP);
  pinMode(SELECT_BUTTON_PIN2, INPUT_PULLUP);
  
  // Initialize OLED
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for(;;);
  }
  
  // Show initial display buffer
  display.display();
  delay(2000);
  
  // Clear the buffer
  display.clearDisplay();
  
  // Draw welcome screen
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.setCursor(10, 20);
  display.println("ESP32 OLED");
  display.setCursor(25, 40);
  display.println("Demo App");
  display.display();
  delay(2500);
  
  // Initialize built-in LED
  pinMode(LED_BUILTIN, OUTPUT);
  digitalWrite(LED_BUILTIN, LOW);
}

void loop() {
  // Read button states
  bool upPressed = (digitalRead(UP_BUTTON_PIN1) == LOW || digitalRead(UP_BUTTON_PIN2) == LOW);
  bool downPressed = (digitalRead(DOWN_BUTTON_PIN1) == LOW || digitalRead(DOWN_BUTTON_PIN2) == LOW);
  bool leftPressed = (digitalRead(LEFT_BUTTON_PIN1) == LOW || digitalRead(LEFT_BUTTON_PIN2) == LOW);
  bool rightPressed = (digitalRead(RIGHT_BUTTON_PIN1) == LOW || digitalRead(RIGHT_BUTTON_PIN2) == LOW);
  bool selectPressed = (digitalRead(SELECT_BUTTON_PIN1) == LOW || digitalRead(SELECT_BUTTON_PIN2) == LOW);
  
  // Handle button presses based on current state
  switch(currentState) {
    case MAIN_MENU:
      handleMainMenu(upPressed, downPressed, selectPressed);
      break;
    case SENSOR_DISPLAY:
      handleSensorDisplay(selectPressed);
      break;
    case LED_CONTROL:
      handleLEDControl(upPressed, downPressed, leftPressed, rightPressed, selectPressed);
      break;
    case SETTINGS:
      handleSettings(upPressed, downPressed, selectPressed);
      break;
  }
  
  // Small delay to debounce buttons
  delay(50);
}

void handleMainMenu(bool upPressed, bool downPressed, bool selectPressed) {
  // Navigation
  if (upPressed) {
    menuIndex = (menuIndex - 1 + 4) % 4;
    drawMainMenu();
  }
  if (downPressed) {
    menuIndex = (menuIndex + 1) % 4;
    drawMainMenu();
  }
  
  // Selection
  if (selectPressed) {
    switch(menuIndex) {
      case 0: 
        currentState = SENSOR_DISPLAY;
        drawSensorDisplay();
        break;
      case 1: 
        currentState = LED_CONTROL;
        drawLEDControl();
        break;
      case 2: 
        currentState = SETTINGS;
        drawSettings();
        break;
      case 3: 
        drawInfoScreen();
        break;
    }
  }
  
  // Draw menu if not already drawn
  static int lastMenuIndex = -1;
  if (lastMenuIndex != menuIndex) {
    drawMainMenu();
    lastMenuIndex = menuIndex;
  }
}

void handleSensorDisplay(bool selectPressed) {
  // Return to main menu on select press
  if (selectPressed) {
    currentState = MAIN_MENU;
    drawMainMenu();
    return;
  }
  
  // Update sensor values (simulated)
  static unsigned long lastUpdate = 0;
  if (millis() - lastUpdate > 2000) {
    temperature += (random(-10, 10)) / 10.0;
    humidity += (random(-10, 10)) / 10.0;
    
    // Constrain values to reasonable ranges
    temperature = constrain(temperature, 20.0, 30.0);
    humidity = constrain(humidity, 40.0, 60.0);
    
    drawSensorDisplay();
    lastUpdate = millis();
  }
}

void handleLEDControl(bool upPressed, bool downPressed, bool leftPressed, bool rightPressed, bool selectPressed) {
  // Toggle LED state with up/down
  if (upPressed) {
    ledState = true;
    digitalWrite(LED_BUILTIN, HIGH);
    drawLEDControl();
  }
  if (downPressed) {
    ledState = false;
    digitalWrite(LED_BUILTIN, LOW);
    drawLEDControl();
  }
  
  // Adjust brightness with left/right
  if (leftPressed) {
    ledBrightness = constrain(ledBrightness - 10, 0, 255);
    if (ledState) analogWrite(LED_BUILTIN, ledBrightness);
    drawLEDControl();
  }
  if (rightPressed) {
    ledBrightness = constrain(ledBrightness + 10, 0, 255);
    if (ledState) analogWrite(LED_BUILTIN, ledBrightness);
    drawLEDControl();
  }
  
  // Return to main menu on select press
  if (selectPressed) {
    currentState = MAIN_MENU;
    drawMainMenu();
  }
}

void handleSettings(bool upPressed, bool downPressed, bool selectPressed) {
  // Simple settings screen with option to return
  if (selectPressed) {
    currentState = MAIN_MENU;
    drawMainMenu();
  }
}

void drawMainMenu() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  
  // Draw title
  display.setCursor(0, 0);
  display.println("MAIN MENU");
  display.drawLine(0, 10, SCREEN_WIDTH, 10, WHITE);
  
  // Draw menu items
  for (int i = 0; i < 4; i++) {
    if (i == menuIndex) {
      display.setTextColor(BLACK, WHITE); // Inverted text for selected item
      display.print("> ");
    } else {
      display.setTextColor(WHITE);
      display.print("  ");
    }
    display.println(menuItems[i]);
  }
  
  // Draw instructions
  display.setTextColor(WHITE);
  display.setCursor(0, SCREEN_HEIGHT - 8);
  display.println("UP/DOWN:Nav  SEL:Select");
  
  display.display();
}

void drawSensorDisplay() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  
  // Draw title
  display.setCursor(0, 0);
  display.println("SENSOR DATA");
  display.drawLine(0, 10, SCREEN_WIDTH, 10, WHITE);
  
  // Draw temperature
  display.setCursor(0, 15);
  display.print("Temperature: ");
  display.print(temperature, 1);
  display.println(" C");
  
  // Draw humidity
  display.setCursor(0, 25);
  display.print("Humidity: ");
  display.print(humidity, 1);
  display.println(" %");
  
  // Draw simulated graph
  display.drawRect(0, 40, SCREEN_WIDTH, 20, WHITE);
  int tempBar = map(temperature, 20, 30, 0, SCREEN_WIDTH-2);
  int humBar = map(humidity, 40, 60, 0, SCREEN_WIDTH-2);
  
  display.fillRect(1, 41, tempBar, 8, WHITE);
  display.fillRect(1, 51, humBar, 8, WHITE);
  
  display.setCursor(0, 35);
  display.print("T");
  display.setCursor(0, 45);
  display.print("H");
  
  // Draw instructions
  display.setCursor(0, SCREEN_HEIGHT - 8);
  display.println("SEL:Back");
  
  display.display();
}

void drawLEDControl() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  
  // Draw title
  display.setCursor(0, 0);
  display.println("LED CONTROL");
  display.drawLine(0, 10, SCREEN_WIDTH, 10, WHITE);
  
  // Draw LED state
  display.setCursor(0, 15);
  display.print("State: ");
  display.println(ledState ? "ON" : "OFF");
  
  // Draw brightness
  display.setCursor(0, 25);
  display.print("Brightness: ");
  display.println(ledBrightness);
  
  // Draw brightness bar
  display.drawRect(0, 35, SCREEN_WIDTH, 10, WHITE);
  int barWidth = map(ledBrightness, 0, 255, 0, SCREEN_WIDTH-2);
  display.fillRect(1, 36, barWidth, 8, WHITE);
  
  // Draw LED symbol
  display.drawCircle(SCREEN_WIDTH/2, 55, 10, WHITE);
  if (ledState) {
    display.fillCircle(SCREEN_WIDTH/2, 55, 8, WHITE);
  }
  
  // Draw instructions
  display.setCursor(0, SCREEN_HEIGHT - 8);
  display.println("UP/DOWN:State L/R:Brght SEL:Back");
  
  display.display();
}

void drawSettings() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  
  // Draw title
  display.setCursor(0, 0);
  display.println("SETTINGS");
  display.drawLine(0, 10, SCREEN_WIDTH, 10, WHITE);
  
  // Draw settings content
  display.setCursor(0, 15);
  display.println("Firmware Version: 1.0");
  display.setCursor(0, 25);
  display.println("OLED: SSD1306 128x64");
  display.setCursor(0, 35);
  display.println("ESP32 Dev Module");
  
  // Draw instructions
  display.setCursor(0, SCREEN_HEIGHT - 8);
  display.println("SEL:Back");
  
  display.display();
}

void drawInfoScreen() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  
  // Draw title
  display.setCursor(0, 0);
  display.println("SYSTEM INFO");
  display.drawLine(0, 10, SCREEN_WIDTH, 10, WHITE);
  
  // Draw info content
  display.setCursor(0, 15);
  display.print("Chip: ");
  display.println(ESP.getChipModel());
  display.setCursor(0, 25);
  display.print("CPU Freq: ");
  display.print(ESP.getCpuFreqMHz());
  display.println(" MHz");
  display.setCursor(0, 35);
  display.print("Flash: ");
  display.print(ESP.getFlashChipSize() / (1024 * 1024));
  display.println(" MB");
  display.setCursor(0, 45);
  display.print("Free RAM: ");
  display.print(ESP.getFreeHeap() / 1024);
  display.println(" KB");
  
  // Draw instructions
  display.setCursor(0, SCREEN_HEIGHT - 8);
  display.println("SEL:Back");
  
  display.display();
  
  // Wait for select button to return
  while (true) {
    if (digitalRead(SELECT_BUTTON_PIN1) == LOW || digitalRead(SELECT_BUTTON_PIN2) == LOW) {
      currentState = MAIN_MENU;
      drawMainMenu();
      break;
    }
    delay(50);
  }
}