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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 32
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// === BUTTON PINS ===
#define BTN_DOWN 8
#define BTN_SELECT 9
#define BTN_UP 10

// === MENU STATE ===
enum State { MAIN_MENU, SUB_MENU_A, SUB_MENU_B, EMERGENCY_STOP };
State currentState = MAIN_MENU;

int mainIndex = 0;
int subIndex = 0;
int scrollOffset = 0;

unsigned long pressTime = 0;
bool estop = false;

// === MENU ARRAYS ===
String mainMenu[] = {"Motor A", "Motor B", "STOP"};
const int mainCount = 3;

String subMenu[] = {"Speed", "Rotation", "Stop", "Return"};
const int subCount = 4;

// === MOTOR STATES ===
int speedA = 0;
int speedB = 0;
bool dirA = true;
bool dirB = true;

void setup() {
  Serial.begin(115200);
  delay(200);
  pinMode(BTN_DOWN, INPUT_PULLDOWN);
  pinMode(BTN_SELECT, INPUT_PULLDOWN);
  pinMode(BTN_UP, INPUT_PULLDOWN);

  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println("Display not found!");
    while(1);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 10);
  display.println("OLED & Menu ready");
  display.display();
  delay(1000);
}

void loop() {
  // === EMERGENCY STOP COMBO ===
  if (digitalRead(BTN_UP) && digitalRead(BTN_DOWN)) {
    if (millis() - pressTime > 1500) {
      estop = !estop;
      if (estop) currentState = EMERGENCY_STOP;
      else currentState = MAIN_MENU;
      delay(500);
    }
  } else {
    pressTime = millis();
  }

  if (estop) {
    drawEstop();
    return;
  }

  // === INPUT HANDLING ===
  if (digitalRead(BTN_UP)) {
    if (currentState == MAIN_MENU) {
      mainIndex--;
      if (mainIndex < 0) mainIndex = mainCount - 1;
    } else if (currentState == SUB_MENU_A || currentState == SUB_MENU_B) {
      subIndex--;
      if (subIndex < 0) subIndex = subCount - 1;
    }
    delay(200);
  }

  if (digitalRead(BTN_DOWN)) {
    if (currentState == MAIN_MENU) {
      mainIndex++;
      if (mainIndex >= mainCount) mainIndex = 0;
    } else if (currentState == SUB_MENU_A || currentState == SUB_MENU_B) {
      subIndex++;
      if (subIndex >= subCount) subIndex = 0;
    }
    delay(200);
  }

  if (digitalRead(BTN_SELECT)) {
    handleSelect();
    delay(300);
  }

  // === DRAW UI ===
  if (currentState == MAIN_MENU)
    drawMenu(mainMenu, mainCount, mainIndex, "Main Menu");
  else if (currentState == SUB_MENU_A)
    drawMenu(subMenu, subCount, subIndex, "Motor A Menu");
  else if (currentState == SUB_MENU_B)
    drawMenu(subMenu, subCount, subIndex, "Motor B Menu");
}

void handleSelect() {
  if (currentState == MAIN_MENU) {
    if (mainIndex == 0) { currentState = SUB_MENU_A; subIndex = 0; }
    else if (mainIndex == 1) { currentState = SUB_MENU_B; subIndex = 0; }
    else if (mainIndex == 2) { // STOP ALL
      speedA = 0;
      speedB = 0;
    }
  } else if (currentState == SUB_MENU_A) {
    performAction('A');
  } else if (currentState == SUB_MENU_B) {
    performAction('B');
  }
}

void performAction(char motor) {
  switch(subIndex) {
    case 0: // Speed
      if (motor == 'A') speedA += 10;
      else speedB += 10;
      if (speedA > 255) speedA = 0;
      if (speedB > 255) speedB = 0;
      break;
    case 1: // Rotation
      if (motor == 'A') dirA = !dirA;
      else dirB = !dirB;
      break;
    case 2: // Stop
      if (motor == 'A') speedA = 0;
      else speedB = 0;
      break;
    case 3: // Return
      currentState = MAIN_MENU;
      break;
  }
}

void drawMenu(String menu[], int count, int index, String title) {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  // Scroll offset handling
  int visibleLines = 3; // 4 lines total - 1 reserved for status bar
  if (index < scrollOffset) scrollOffset = index;
  else if (index >= scrollOffset + visibleLines) scrollOffset = index - visibleLines + 1;

  display.setCursor(0, 0);
  display.println(title);

  for (int i = scrollOffset; i < min(count, scrollOffset + visibleLines); i++) {
    if (i == index) display.print("> ");
    else display.print("  ");
    display.println(menu[i]);
  }

  // Bottom status bar
  display.drawLine(0, 24, 128, 24, SSD1306_WHITE);
  display.setCursor(0, 26);
  display.setTextSize(1);
  display.printf("A:%03d%c  B:%03d%c",
    speedA, dirA ? '>' : '<',
    speedB, dirB ? '>' : '<');

  display.display();
}

void drawEstop() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 10);
  display.println("!!! EMERGENCY STOP !!!");
  display.display();
}