#include <Adafruit_ILI9341.h>

#define TFT_CS   15
#define TFT_DC   2
#define TFT_WIDTH 320
#define TFT_HEIGHT 240
Adafruit_ILI9341 tft = Adafruit_ILI9341(TFT_CS, TFT_DC);

#define SOLO_MODE 0
#define MULTI_MODE 1
#define SETTINGS_MODE 2
#define CHARACTER_SELECTION_MODE 3
#define LEVEL_SELECTION_MODE 4

int currentMode = SOLO_MODE;
int selectedOption = SOLO_MODE;
unsigned long lastButtonPressTime = 0;

int selectedCharacter = 0;
int selectedLevel = 0;

#define ILI9341_PINK 0xFC9F
#define TFT_WHITE   0xFFFF
#define TFT_RED     0xF800
#define TFT_BLUE    0x001F
#define TFT_GREEN   0x07E0

#define TITLE_COLOR TFT_BLUE
#define SELECTED_COLOR TFT_RED
#define UNSELECTED_COLOR TFT_WHITE
#define ARROW_COLOR TFT_GREEN

#define TITLE_TEXT_SIZE 2
#define OPTION_TEXT_SIZE 2
#define CHARACTER_SIZE 20
#define CHARACTER_SPACING 60
#define OPTION_Y_START (TFT_HEIGHT / 2)

#define BTN_UP    16
#define BTN_DOWN  17
#define BTN_SELECT 5
#define BTN_LEFT  25  // Left button for character selection
#define BTN_RIGHT 26  // Right button for character selection

void setup() {
  Serial.begin(9600);
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_SELECT, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);
  
  tft.begin();
  tft.setRotation(1);
  tft.fillScreen(ILI9341_PINK);
  
  displayHomePage();
}

void loop() {
  handleButtonPresses();
}

void displayHomePage() {
  tft.fillScreen(ILI9341_PINK);
  tft.setTextSize(TITLE_TEXT_SIZE);
  tft.setTextColor(TITLE_COLOR);
  tft.setCursor((TFT_WIDTH / 2) - 100, 30);
  tft.println("The Legend of Shit");

  currentMode = SOLO_MODE;
  displayOptions();
}

void displayOptions() {
  tft.setTextSize(OPTION_TEXT_SIZE);
  int optionY = OPTION_Y_START - 60;  // Start position for options

  for (int i = 0; i < 3; i++) {
    tft.setTextColor(i == selectedOption ? SELECTED_COLOR : UNSELECTED_COLOR);
    tft.setCursor((TFT_WIDTH / 2) - 30, optionY + 50);
    switch (i) {
      case SOLO_MODE:
        tft.println("Solo");
        break;
      case MULTI_MODE:
        tft.println("Multi");
        break;
      case SETTINGS_MODE:
        tft.println("Setting");
        break;
    }
    optionY += 30; // Increase Y for next option
  }
}

void displayCharacterSelectionPage() {
  tft.fillScreen(ILI9341_PINK);
  tft.setTextSize(TITLE_TEXT_SIZE);
  tft.setTextColor(TITLE_COLOR);
  tft.setCursor((TFT_WIDTH / 2) - 90, 20);
  tft.println("Select Character");

  for (int i = 0; i < 3; i++) {
    updateCharacterSelection(i, i == selectedCharacter ? SELECTED_COLOR : UNSELECTED_COLOR);
  }
  displayArrowIndicators();
}

void updateCharacterSelection(int character, uint16_t color) {
  int firstCharacterX = (TFT_WIDTH / 2) - (CHARACTER_SPACING * 1.5);
  int characterX = firstCharacterX + character * CHARACTER_SPACING;
  int characterY = OPTION_Y_START + 40;

  tft.fillCircle(characterX + 20, characterY - 20, CHARACTER_SIZE / 2, color);
  tft.setTextSize(OPTION_TEXT_SIZE);
  tft.setTextColor(color);
  tft.setCursor(characterX + 10, characterY + 20);
  switch (character) {
    case 0:
      tft.println("Ray");
      break;
    case 1:
      tft.println("Max");
      break;
    case 2:
      tft.println("Zoe");
      break;
  }
}

void displayArrowIndicators() {
  int characterY = OPTION_Y_START + 40;
  int arrowY = characterY + 20;
  int arrowLeftX = (TFT_WIDTH / 2) - (CHARACTER_SPACING * 1.5) - 50;
  int arrowRightX = (TFT_WIDTH / 2) + (CHARACTER_SPACING * 1.5) + 30;

  tft.setTextSize(OPTION_TEXT_SIZE);
  tft.setTextColor(ARROW_COLOR);
  tft.setCursor(arrowLeftX, arrowY);
  tft.print("<");
  tft.setCursor(arrowRightX, arrowY);
  tft.print(">");
}

void displayLevelSelectionPage() {
  tft.fillScreen(ILI9341_PINK);
  tft.setTextSize(TITLE_TEXT_SIZE);
  tft.setTextColor(TITLE_COLOR);
  tft.setCursor((TFT_WIDTH / 2) - 50, 20);
  tft.println("Select Level");

  int levelY = OPTION_Y_START - 60;  // Start position for levels
  String levels[] = {"Easy", "Medium", "Hard"};

  for (int i = 0; i < 3; i++) {
    displayLevelOption(i, i == selectedLevel);
  }
}

void displayLevelOption(int levelIndex, bool isSelected) {
  int levelY = OPTION_Y_START + (levelIndex * 30) - 60; // Calculate Y position for level
  tft.setTextColor(isSelected ? SELECTED_COLOR : UNSELECTED_COLOR);
  tft.setCursor((TFT_WIDTH / 2) - 30, levelY);
  switch (levelIndex) {
    case 0:
      tft.println("Easy   "); // Extra spaces to clear any previous longer text
      break;
    case 1:
      tft.println("Medium ");
      break;
    case 2:
      tft.println("Hard   ");
      break;
  }
}

void updateLevelSelection(int oldLevel, int newLevel) {
  displayLevelOption(oldLevel, false); // Redraw old option as unselected
  displayLevelOption(newLevel, true);  // Redraw new option as selected
}

void handleButtonPresses() {
  unsigned long currentMillis = millis();
  if ((currentMode == SOLO_MODE || currentMode == LEVEL_SELECTION_MODE) && (digitalRead(BTN_UP) == LOW || digitalRead(BTN_DOWN) == LOW)) {
    if (currentMillis - lastButtonPressTime > 200) {
      lastButtonPressTime = currentMillis;
      int oldLevel = selectedLevel;
      if (digitalRead(BTN_UP) == LOW) {
        if (currentMode == SOLO_MODE && selectedOption > 0) {
          selectedOption--;
          displayOptions();
        }
        if (currentMode == LEVEL_SELECTION_MODE && selectedLevel > 0) {
          selectedLevel--;
          updateLevelSelection(oldLevel, selectedLevel);
        }
      }
      if (digitalRead(BTN_DOWN) == LOW) {
        if (currentMode == SOLO_MODE && selectedOption < 2) {
          selectedOption++;
          displayOptions();
        }
        if (currentMode == LEVEL_SELECTION_MODE && selectedLevel < 2) {
          selectedLevel++;
          updateLevelSelection(oldLevel, selectedLevel);
        }
      }
    }
  } else if (currentMode == CHARACTER_SELECTION_MODE) {
    if (digitalRead(BTN_LEFT) == LOW && currentMillis - lastButtonPressTime > 200) {
      lastButtonPressTime = currentMillis;
      if (selectedCharacter > 0) {
        selectedCharacter--;
        updateCharacterSelection(selectedCharacter + 1, UNSELECTED_COLOR);
        updateCharacterSelection(selectedCharacter, SELECTED_COLOR);
        displayArrowIndicators();
      }
    }
    if (digitalRead(BTN_RIGHT) == LOW && currentMillis - lastButtonPressTime > 200) {
      lastButtonPressTime = currentMillis;
      if (selectedCharacter < 2) {
        selectedCharacter++;
        updateCharacterSelection(selectedCharacter - 1, UNSELECTED_COLOR);
        updateCharacterSelection(selectedCharacter, SELECTED_COLOR);
        displayArrowIndicators();
      }
    }
  }
  
  if (digitalRead(BTN_SELECT) == LOW && currentMillis - lastButtonPressTime > 200) {
    lastButtonPressTime = currentMillis;
    if (currentMode == SOLO_MODE) {
      currentMode = CHARACTER_SELECTION_MODE;
      selectedCharacter = 0;
      displayCharacterSelectionPage();
    } else if (currentMode == CHARACTER_SELECTION_MODE) {
      currentMode = LEVEL_SELECTION_MODE;
      selectedLevel = 0;
      displayLevelSelectionPage();
    }
  }
}