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

#define SCREEN_WIDTH 128 // OLED display width, in pixels
#define SCREEN_HEIGHT 64 // OLED display height, in pixels

#define OLED_RESET -1 // Reset pin # (or -1 if sharing Arduino reset pin)
#define SCREEN_ADDRESS 0x3C // Address for 128x64 screen
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// control buttons
#define button_Y 0
#define button_B 1
#define button_A 2
#define button_X 3
#define nav_switch_up 4
#define nav_switch_select 5
#define nav_switch_down 6

int MenuItem = 0;
int ItemsPerPage = 3;

// Menu tracking variables
int currentMenu = 0; // 0 for main menu, 1 for games menu, etc.

// Define menu items
const char* mainMenuItems[] = {"Settings", "Games", "Music"};
const char* gamesMenuItems[] = {"FlappyBird", "StarWars", "Sudoku"};
const int mainMenuSize = 3;
const int gamesMenuSize = 3;

// Button state tracking variables
bool lastNavSwitchDownState = HIGH;
bool lastNavSwitchUpState = HIGH;
bool lastNavSwitchSelectState = HIGH;

void setup() {
  Serial.begin(9600);

  // Initialize the OLED display
  if (!display.begin(SSD1306_SWITCHCAPVCC, SCREEN_ADDRESS)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;); // Don't proceed, loop forever
  }

  // clear display
  display.clearDisplay();
  display.setTextColor(WHITE);

  // input buttons
  pinMode(nav_switch_up, INPUT_PULLUP);
  pinMode(nav_switch_select, INPUT_PULLUP);
  pinMode(nav_switch_down, INPUT_PULLUP);

  // show main menu
  displayMenu();
}

void loop() {
  // Read the current state of the buttons
  bool currentNavSwitchDownState = digitalRead(nav_switch_down);
  bool currentNavSwitchUpState = digitalRead(nav_switch_up);
  bool currentNavSwitchSelectState = digitalRead(nav_switch_select);

  // Handle down navigation
  if (lastNavSwitchDownState == HIGH && currentNavSwitchDownState == LOW) {
    MenuItem++;
    int menuSize = (currentMenu == 0) ? mainMenuSize : gamesMenuSize;
    if (MenuItem >= menuSize) {
      MenuItem = 0; // Go back to the first item
    }
    displayMenu();
  }

  // Handle up navigation
  if (lastNavSwitchUpState == HIGH && currentNavSwitchUpState == LOW) {
    MenuItem--;
    int menuSize = (currentMenu == 0) ? mainMenuSize : gamesMenuSize;
    if (MenuItem < 0) {
      MenuItem = menuSize - 1; // Go to the last item
    }
    displayMenu();
  }

  // Handle select button (open submenu)
  if (lastNavSwitchSelectState == HIGH && currentNavSwitchSelectState == LOW) {
    if (currentMenu == 0) { // Main menu
      if (MenuItem == 1) {
        // User selected "Games", open games menu
        currentMenu = 1;
        MenuItem = 0; // Reset to first item of games menu
      }
    }
  }

  // Update the last button states for the next loop
  lastNavSwitchDownState = currentNavSwitchDownState;
  lastNavSwitchUpState = currentNavSwitchUpState;
  lastNavSwitchSelectState = currentNavSwitchSelectState;
}

// Function to display the current menu
void displayMenu() {
  display.clearDisplay();

  const char** menuItems;
  int menuSize;

  if (currentMenu == 0) {
    // Main Menu
    menuItems = mainMenuItems;
    menuSize = mainMenuSize;
  } else if (currentMenu == 1) {
    // Games Menu
    menuItems = gamesMenuItems;
    menuSize = gamesMenuSize;
  }

  // Display menu items for the current menu
  int FirstItemOfPage = (MenuItem / ItemsPerPage) * ItemsPerPage;
  for (int i = 0; i < ItemsPerPage; i++) {
    int CurrentItem = FirstItemOfPage + i;

    if (CurrentItem < menuSize) {
      display.setCursor(10, 10 + (i * 15));

      if (CurrentItem == MenuItem) {
        display.print("> ");
      } else {
        display.print("  ");
      }

      display.print(menuItems[CurrentItem]);
    }
  }

  display.display();
}