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

#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
#define EEPROM_SIZE 10

// Initialize the display
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Button pins
const int BUTTON_UP_PIN = 15;
const int BUTTON_DOWN_PIN = 4;
const int BUTTON_LEFT_PIN = 14;
const int BUTTON_RIGHT_PIN = 13;
const int BUTTON_SELECT_PIN = 27;

// Menu and game states
enum GameState {
  MENU,
  PONG,
  CATCH_THE_BALL
};

GameState currentState = MENU;
int menuSelection = 0;

// Pong game properties
const int PADDLE_WIDTH = 8;
const int PADDLE_HEIGHT = 16;
const int BALL_SIZE = 4;
int paddleY = (SCREEN_HEIGHT - PADDLE_HEIGHT) / 2;
int ballX = SCREEN_WIDTH / 2;
int ballY = SCREEN_HEIGHT / 2;
int ballSpeedX = 2;
int ballSpeedY = 2;
int pongScore = 0; // Score for Pong game

// Catch the Ball game properties
const int CATCH_PADDLE_WIDTH = 16;
const int CATCH_PADDLE_HEIGHT = 4;
int catchPaddleX = (SCREEN_WIDTH - CATCH_PADDLE_WIDTH) / 2;
int catchPaddleY = SCREEN_HEIGHT - CATCH_PADDLE_HEIGHT;
int fallingBallX = random(0, SCREEN_WIDTH - BALL_SIZE);
int fallingBallY = 0;
int fallingBallSpeedY = 2;
int catchScore = 0; 
int highScore = 0; 

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

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  pinMode(BUTTON_UP_PIN, INPUT_PULLUP);
  pinMode(BUTTON_DOWN_PIN, INPUT_PULLUP);
  pinMode(BUTTON_LEFT_PIN, INPUT_PULLUP);
  pinMode(BUTTON_RIGHT_PIN, INPUT_PULLUP);
  pinMode(BUTTON_SELECT_PIN, INPUT_PULLUP);

  // Read high score from EEPROM
  EEPROM.begin(EEPROM_SIZE);
  highScore = EEPROM.read(0);

  displayMenu();
}

void loop() {
  int readingUp = digitalRead(BUTTON_UP_PIN);
  int readingDown = digitalRead(BUTTON_DOWN_PIN);
  int readingLeft = digitalRead(BUTTON_LEFT_PIN);
  int readingRight = digitalRead(BUTTON_RIGHT_PIN);
  int readingSelect = digitalRead(BUTTON_SELECT_PIN);

  // Button debounce handling
  if (readingUp == LOW) {
    delay(50); // Simple debounce delay
    if (currentState == MENU) {
      menuSelection = (menuSelection - 1 + 2) % 2; // Toggle between 0 and 1
      displayMenu();
    }
  }

  if (readingDown == LOW) {
    delay(50); // Simple debounce delay
    if (currentState == MENU) {
      menuSelection = (menuSelection + 1) % 2; // Toggle between 0 and 1
      displayMenu();
    }
  }

  if (readingSelect == LOW) {
    delay(50); // Simple debounce delay
    if (currentState == MENU) {
      if (menuSelection == 0) {
        currentState = PONG;
        playPong();
      } else if (menuSelection == 1) {
        currentState = CATCH_THE_BALL;
        playCatchTheBall();
      }
    }
  }
}

void displayMenu() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 0);
  display.println("Game Menu");
  display.setTextSize(1);
  if (menuSelection == 0) {
    display.setCursor(0, 30);
    display.println("> Pong");
    display.setCursor(0, 40);
    display.println("  Catch the Ball");
  } else {
    display.setCursor(0, 30);
    display.println("  Pong");
    display.setCursor(0, 40);
    display.println("> Catch the Ball");
  }
  display.display();
}

void playPong() {
  paddleY = (SCREEN_HEIGHT - PADDLE_HEIGHT) / 2;
  ballX = SCREEN_WIDTH / 2;
  ballY = SCREEN_HEIGHT / 2;
  ballSpeedX = 2;
  ballSpeedY = 2;
  pongScore = 0; // Reset score for Pong game

  while (currentState == PONG) {
    int readingUp = digitalRead(BUTTON_UP_PIN);
    int readingDown = digitalRead(BUTTON_DOWN_PIN);

    // Handle button press for UP
    if (readingUp == LOW) {
      if (paddleY > 0) {
        paddleY -= 2;
      }
      delay(50); // Simple debounce delay
    }

    // Handle button press for DOWN
    if (readingDown == LOW) {
      if (paddleY < SCREEN_HEIGHT - PADDLE_HEIGHT) {
        paddleY += 2;
      }
      delay(50); // Simple debounce delay
    }

    // Move the ball
    ballX += ballSpeedX;
    ballY += ballSpeedY;

    // Ball collision with top and bottom edges
    if (ballY <= 0 || ballY >= SCREEN_HEIGHT - BALL_SIZE) {
      ballSpeedY = -ballSpeedY;
    }

    // Ball collision with the paddle
    if (ballX <= PADDLE_WIDTH && ballY >= paddleY && ballY <= paddleY + PADDLE_HEIGHT) {
      ballSpeedX = -ballSpeedX;
      ballX = PADDLE_WIDTH; // Ensure the ball doesn't get stuck in the paddle
      pongScore++; // Increase score when ball hits the paddle
    }

    // Ball collision with the right edge (rebound instead of respawn)
    if (ballX >= SCREEN_WIDTH - BALL_SIZE) {
      ballSpeedX = -ballSpeedX; // Reverse the direction
      ballX = SCREEN_WIDTH - BALL_SIZE; // Ensure the ball is within screen bounds
      pongScore++; // Increase score when ball hits the right edge
    }

    // Clear the display
    display.clearDisplay();

    // Draw the paddle
    display.fillRect(0, paddleY, PADDLE_WIDTH, PADDLE_HEIGHT, SSD1306_WHITE);

    // Draw the ball
    display.fillRect(ballX, ballY, BALL_SIZE, BALL_SIZE, SSD1306_WHITE);

    // Draw the score
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.print("Score: ");
    display.println(pongScore);

    // Update the display
    display.display();

    // Delay for a short period to control the game speed
    delay(20);

    // Check if SELECT button is pressed to return to menu
    if (digitalRead(BUTTON_SELECT_PIN) == LOW) {
      currentState = MENU;
      displayMenu();
      return; // Exit the function to prevent continuing the game loop
    }
  }

  // After game ends, display final score and update high score
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 0);
  display.println("Game Over!");
  display.setTextSize(1);
  display.setCursor(0, 30);
  display.print("Final Score: ");
  display.println(pongScore);
  
  // Check and update high score
  if (pongScore > highScore) {
    highScore = pongScore;
    EEPROM.write(0, highScore);
    EEPROM.commit(); // Ensure the new high score is saved
  }

  display.setCursor(0, 50);
  display.print("High Score: ");
  display.println(highScore);
  
  display.display();
  
  // Wait for SELECT button to return to menu
  while (digitalRead(BUTTON_SELECT_PIN) == HIGH) {
    delay(10);
  }
  
  // Go back to menu
  currentState = MENU;
  displayMenu();
}

void playCatchTheBall() {
  catchPaddleX = (SCREEN_WIDTH - CATCH_PADDLE_WIDTH) / 2;
  catchPaddleY = SCREEN_HEIGHT - CATCH_PADDLE_HEIGHT;
  fallingBallX = random(0, SCREEN_WIDTH - BALL_SIZE);
  fallingBallY = 0;
  catchScore = 0; // Reset score for Catch the Ball game

  while (currentState == CATCH_THE_BALL) {
    int readingUp = digitalRead(BUTTON_UP_PIN);
    int readingDown = digitalRead(BUTTON_DOWN_PIN);
    int readingLeft = digitalRead(BUTTON_LEFT_PIN);
    int readingRight = digitalRead(BUTTON_RIGHT_PIN);

    // Handle paddle movement
    if (readingLeft == LOW) {
      if (catchPaddleX > 0) {
        catchPaddleX -= 2;
      }
      delay(50); // Simple debounce delay
    }

    if (readingRight == LOW) {
      if (catchPaddleX < SCREEN_WIDTH - CATCH_PADDLE_WIDTH) {
        catchPaddleX += 2;
      }
      delay(50); // Simple debounce delay
    }

    if (readingUp == LOW) {
      if (catchPaddleY > 0) {
        catchPaddleY -= 2;
      }
      delay(50); // Simple debounce delay
    }

    if (readingDown == LOW) {
      if (catchPaddleY < SCREEN_HEIGHT - CATCH_PADDLE_HEIGHT) {
        catchPaddleY += 2;
      }
      delay(50); // Simple debounce delay
    }

    // Move the falling ball
    fallingBallY += fallingBallSpeedY;

    // Check if ball falls off screen
    if (fallingBallY >= SCREEN_HEIGHT) {
      fallingBallX = random(0, SCREEN_WIDTH - BALL_SIZE);
      fallingBallY = 0;
    }

    // Ball collision with the paddle
    if (fallingBallY >= catchPaddleY && fallingBallY <= catchPaddleY + CATCH_PADDLE_HEIGHT &&
        fallingBallX >= catchPaddleX && fallingBallX <= catchPaddleX + CATCH_PADDLE_WIDTH) {
      // Increase score when ball is caught
      catchScore++;
      fallingBallX = random(0, SCREEN_WIDTH - BALL_SIZE);
      fallingBallY = 0;
    }

    // Clear the display
    display.clearDisplay();

    // Draw the paddle
    display.fillRect(catchPaddleX, catchPaddleY, CATCH_PADDLE_WIDTH, CATCH_PADDLE_HEIGHT, SSD1306_WHITE);

    // Draw the falling ball
    display.fillRect(fallingBallX, fallingBallY, BALL_SIZE, BALL_SIZE, SSD1306_WHITE);

    // Draw the score
    display.setTextSize(1);
    display.setCursor(0, 0);
    display.print("Score: ");
    display.println(catchScore);

    // Update the display
    display.display();

    // Delay for a short period to control the game speed
    delay(20);

    // Check if SELECT button is pressed to return to menu
    if (digitalRead(BUTTON_SELECT_PIN) == LOW) {
      currentState = MENU;
      displayMenu();
    }
  }

  // After game ends, display final score and update high score
  display.clearDisplay();
  display.setTextSize(2);
  display.setCursor(0, 0);
  display.println("Game Over!");
  display.setTextSize(1);
  display.setCursor(0, 30);
  display.print("Final Score: ");
  display.println(catchScore);
  
  // Check and update high score
  if (catchScore > highScore) {
    highScore = catchScore;
    EEPROM.write(0, highScore);
  }

  display.setCursor(0, 50);
  display.print("High Score: ");
  display.println(highScore);
  
  display.display();
  
  // Wait for SELECT button to return to menu
  while (digitalRead(BUTTON_SELECT_PIN) == HIGH) {
    delay(50);
  }
}