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

// Constants
const int Screen_width = 128;
const int Screen_ht = 64;
const int paddle_width = 20;
const int paddle_height = 4;
const int ball_radius = 2;

// Ball position and speed
int ball_x = 64, ball_y = 32;
int ball_x_speed = 1, ball_y_speed = 1;

// Paddle positions
int paddle1_x = 55, paddle1_y = 60; // Player 1 (bottom)
int paddle2_x = 55, paddle2_y = 0;  // Player 2 (top)

// Scores
int score1 = 0, score2 = 0;

// Button states
int player1_left, player1_right, player2_left, player2_right;

// Game modes
enum GameMode { SINGLE_PLAYER, DUAL_PLAYER, PLAYER_VS_AI };
GameMode game_mode = SINGLE_PLAYER;

// Restart button pin
const int restart_button_pin = 17;

// Ball collision flags
bool ball_hit_paddle1 = false;
bool ball_hit_paddle2 = false;

// Initialize display
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);

void setup() {
  pinMode(16, INPUT_PULLDOWN); // Player 1 left
  pinMode(18, INPUT_PULLDOWN); // Player 1 right
  pinMode(4, INPUT_PULLDOWN);  // Player 2 left
  pinMode(5, INPUT_PULLDOWN);  // Player 2 right
  pinMode(19, OUTPUT);         // Game over LED
  pinMode(restart_button_pin, INPUT_PULLDOWN); // Restart button
  Serial.begin(9600);
  u8g2.begin();

  // Mode selection
  selectMode();
}

void loop() {
  // Check for restart button press
  if (digitalRead(restart_button_pin) == HIGH) {
    resetGame();  // Reset game state when restart button is pressed
    selectMode(); // Ask the user to select the game mode again
  }

  // Display scores
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB08_tr);
  u8g2.setCursor(10, 10);
  u8g2.print("P1: ");
  u8g2.print(score1);

  if (game_mode == DUAL_PLAYER) { // Show Player 2 score only in relevant modes
    u8g2.setCursor(100, 10);
    u8g2.print("P2: ");
    u8g2.print(score2);
  }

  // Paddle 1 (bottom) movement
  player1_right = digitalRead(18);
  if (player1_right && paddle1_x + paddle_width <= 127) {
    paddle1_x += 3;
  }
  player1_left = digitalRead(16);
  if (player1_left && paddle1_x >= 1) {
    paddle1_x -= 3;
  }

  // Paddle 2 movement based on game mode
  if (game_mode == SINGLE_PLAYER || game_mode == PLAYER_VS_AI) {
    // AI control for Paddle 2 (top)
    if (ball_x > paddle2_x + paddle_width / 2 && paddle2_x + paddle_width <= 127) {
      paddle2_x += 2;
    } else if (ball_x < paddle2_x + paddle_width / 2 && paddle2_x >= 1) {
      paddle2_x -= 2;
    }
  } else {
    // Paddle 2 (top) movement for Player 2
    player2_right = digitalRead(5);
    if (player2_right && paddle2_x + paddle_width <= 127) {
      paddle2_x += 3;
    }
    player2_left = digitalRead(4);
    if (player2_left && paddle2_x >= 1) {
      paddle2_x -= 3;
    }
  }

  // Ball boundary conditions
  if (ball_x >= 125 || ball_x <= 2) ball_x_speed = -ball_x_speed;

  // Ball collision with Paddle 1 (bottom)
  if (ball_x >= paddle1_x && ball_x <= paddle1_x + paddle_width && ball_y + ball_radius >= paddle1_y) {
    if (!ball_hit_paddle1 && ball_y_speed > 0) { // Ensure single collision
      ball_y_speed = -ball_y_speed;
      score1++;
      digitalWrite(19, HIGH);
      delay(20);
      digitalWrite(19, LOW);
      ball_hit_paddle1 = true; // Mark collision
    }
  } else {
    ball_hit_paddle1 = false; // Reset collision flag when ball moves away
  }

  // Ball collision with Paddle 2 (top) or upper screen boundary
  if (game_mode == SINGLE_PLAYER || game_mode == PLAYER_VS_AI) {
    // Bounce off the upper screen boundary in AI modes
    if (ball_y - ball_radius <= 0) {
      ball_y_speed = -ball_y_speed;
    }
  } else {
    // Ball collision with Paddle 2 (top) in dual-player mode
    if (ball_x >= paddle2_x && ball_x <= paddle2_x + paddle_width && ball_y - ball_radius <= paddle2_y + paddle_height) {
      if (!ball_hit_paddle2 && ball_y_speed < 0) { // Ensure single collision
        ball_y_speed = -ball_y_speed;
        score2++;
        digitalWrite(19, HIGH);
        delay(20);
        digitalWrite(19, LOW);
        ball_hit_paddle2 = true; // Mark collision
      }
    } else {
      ball_hit_paddle2 = false; // Reset collision flag when ball moves away
    }
  }

  // Update ball position
  ball_x += ball_x_speed;
  ball_y += ball_y_speed;

  // Draw ball and paddles
  u8g2.drawDisc(ball_x, ball_y, ball_radius, U8G2_DRAW_ALL);
  u8g2.drawBox(paddle1_x, paddle1_y, paddle_width, paddle_height); // Player 1 paddle
  if (game_mode != SINGLE_PLAYER) {
    u8g2.drawBox(paddle2_x, paddle2_y, paddle_width, paddle_height); // Player 2 paddle (in relevant modes)
  }
  u8g2.sendBuffer();

  // Game over conditions
  if (ball_y >= Screen_ht) { // Ball passes Paddle 1
    displayFinalScore();
  } else if (ball_y <= 0 && game_mode != SINGLE_PLAYER) { // Ball passes Paddle 2
    displayFinalScore();
  }
}

void selectMode() {
  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB08_tr);
  u8g2.setCursor(10, 20);
  u8g2.print("Select Mode:");
  u8g2.setCursor(10, 40);
  u8g2.print("16: Single Player");
  u8g2.setCursor(10, 50);
  u8g2.print("18: Dual Player");
  u8g2.setCursor(10, 60);
  u8g2.print("4: Player vs AI");
  u8g2.sendBuffer();

  // Wait for mode selection
  bool modeSelected = false;
  while (!modeSelected) {
    if (digitalRead(16) == HIGH) {
      game_mode = SINGLE_PLAYER;
      modeSelected = true;
      Serial.println("Single Player Mode Selected");
    } else if (digitalRead(18) == HIGH) {
      game_mode = DUAL_PLAYER;
      modeSelected = true;
      Serial.println("Dual Player Mode Selected");
    } else if (digitalRead(4) == HIGH) {
      game_mode = PLAYER_VS_AI;
      modeSelected = true;
      Serial.println("Player vs AI Mode Selected");
    }
    delay(50); // Debounce delay
  }

  delay(500); // Prevent immediate re-triggering
}

void displayFinalScore() {
  digitalWrite(19, HIGH);
  delay(2000);
  digitalWrite(19, LOW);

  u8g2.clearBuffer();
  u8g2.setFont(u8g2_font_ncenB08_tr);
  u8g2.setCursor(25, 20);
  u8g2.print("Game Over");
  u8g2.setCursor(10, 40);
  u8g2.print("Final Score P1: ");
  u8g2.print(score1);

  if (game_mode == DUAL_PLAYER) { // Show Player 2's final score in relevant modes
    u8g2.setCursor(10, 50);
    u8g2.print("Final Score P2: ");
    u8g2.print(score2);
  }

  u8g2.sendBuffer();

  // Wait for button press to restart
  waitForRestart();
}

void waitForRestart() {
  bool restart = false;
  while (!restart) {
    if (digitalRead(16) == HIGH || digitalRead(18) == HIGH || digitalRead(4) == HIGH || digitalRead(5) == HIGH || digitalRead(restart_button_pin) == HIGH) {
      restart = true;
    }
    delay(50); // Debounce delay
  }

  resetGame();
}

void resetGame() {
  ball_x = 64;
  ball_y = 32;
  ball_x_speed = 1;
  ball_y_speed = 1;
  score1 = 0;
  score2 = 0;
  paddle1_x = 55;
  paddle2_x = 55;
}