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

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

// Define menu items
const char* menuItems[] = {"Item 1", "Item 2", "Item 3", "Item 4", "Item 5"};
const int menuLength = sizeof(menuItems) / sizeof(menuItems[0]);

// Button pins
const int buttonUpPin = 27;
const int buttonDownPin = 25;
const int buttonSelectPin = 26;

// Variables to keep track of menu state
volatile int selectedItem = 0;
volatile bool inSubMenu = false;
volatile bool buttonUpPressed = false;
volatile bool buttonDownPressed = false;
volatile bool buttonSelectPressed = false;

void IRAM_ATTR handleUpButton() {
  buttonUpPressed = true;
}

void IRAM_ATTR handleDownButton() {
  buttonDownPressed = true;
}

void IRAM_ATTR handleSelectButton() {
  buttonSelectPressed = true;
}

void setup() {
  pinMode(buttonUpPin, INPUT_PULLUP);
  pinMode(buttonDownPin, INPUT_PULLUP);
  pinMode(buttonSelectPin, INPUT_PULLUP);

  attachInterrupt(digitalPinToInterrupt(buttonUpPin), handleUpButton, FALLING);
  attachInterrupt(digitalPinToInterrupt(buttonDownPin), handleDownButton, FALLING);
  attachInterrupt(digitalPinToInterrupt(buttonSelectPin), handleSelectButton, FALLING);

  // Initialize the OLED display
  if (!display.begin(SSD1306_I2C_ADDRESS, OLED_RESET)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }
  display.display();
  delay(2000);
  display.clearDisplay();

  drawMenu();
}

void loop() {
  if (buttonSelectPressed) {
    buttonSelectPressed = false;
    if (!inSubMenu) {
      inSubMenu = true;
      drawSubMenu();
    }
  }

  if (!inSubMenu) {
    if (buttonUpPressed) {
      buttonUpPressed = false;
      selectedItem--;
      if (selectedItem < 0) {
        selectedItem = menuLength - 1;
      }
      drawMenu();
    }

    if (buttonDownPressed) {
      buttonDownPressed = false;
      selectedItem++;
      if (selectedItem >= menuLength) {
        selectedItem = 0;
      }
      drawMenu();
    }
  } else {
    if (buttonDownPressed) {
      buttonDownPressed = false;
      inSubMenu = false;
      drawMenu();
    }
  }
}

void drawMenu() {
  display.clearDisplay();

  int middleIndex = 1;  // The middle item position in the 3-item display (0, 1, 2)
  int displayOffset = selectedItem - middleIndex;

  // Adjust displayOffset for looping
  if (displayOffset < 0) {
    displayOffset += menuLength;
  }

  for (int i = 0; i < 3; i++) {
    int itemIndex = (displayOffset + i) % menuLength;
    if (itemIndex < menuLength) {
      if (itemIndex == selectedItem) {
        // Draw a box around the selected item
        display.setTextColor(SSD1306_BLACK, SSD1306_WHITE);
        display.fillRect(0, i * 16, SCREEN_WIDTH, 16, SSD1306_WHITE);
      } else {
        display.setTextColor(SSD1306_WHITE);
      }
      display.setCursor(0, i * 16);
      display.print(menuItems[itemIndex]);
    }
  }

  display.display();
}

void drawSubMenu() {
  display.clearDisplay();

  // Display the selected item name
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(0, 0);
  display.print(menuItems[selectedItem]);

  // Display the "Return to Menu" option
  display.setCursor(0, SCREEN_HEIGHT - 16);
  display.print("Return to Menu");

  display.display();
}