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

#define PADDLE_WIDTH_MIN 16
#define PADDLE_WIDTH_MAX 48
#define PADDLE_HEIGHT 1
#define BALL_SIZE 1
#define BRICK_ROWS 4
#define BRICK_COLUMNS 8
#define BRICK_WIDTH 14
#define BRICK_HEIGHT 6
#define BRICK_GAP 2

#define JOYSTICK_X A0
#define JOYSTICK_Y A1

Adafruit_SSD1306 display(128, 64, &Wire);

float ball_x, ball_y;
float ball_speed_x = 1.0;
float ball_speed_y = -1.0;
float selected_ball_speed = 1.0;

int paddle_x, paddle_y;
int paddle_width = 16;
int paddle_speed = 3;

bool bricks[BRICK_COLUMNS][BRICK_ROWS];
int brick_colors[BRICK_COLUMNS][BRICK_ROWS];

bool game_over = true;
bool win = false;

int menu_selection = 1;
bool game_started = false;
bool return_to_menu = false;

float angle = 0.0;

void setup() {
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.display();
  delay(2000);
  display.clearDisplay();

  pinMode(JOYSTICK_X, INPUT);
  pinMode(JOYSTICK_Y, INPUT);

  for (int i = 0; i < BRICK_COLUMNS; i++) {
    for (int j = 0; j < BRICK_ROWS; j++) {
      bricks[i][j] = true;
      brick_colors[i][j] = WHITE;
    }
  }
}

void loop() {
  if (!game_started || return_to_menu) {
    displayMenu();
    handleMenu();
  } else {
    if (!game_over) {
      handleGameplay();
    }
    drawGame();
  }

  if (game_over && game_started) {
    display.setTextSize(2);
    display.setTextColor(WHITE);
    display.setCursor(16, 24);
    display.fillRect(0, 0, 128, 64, BLACK);
    if (win) {
      display.print("You win!");
    } else {
      display.print("You lose!");
    }
    
    display.display();
    delay(2000);
    resetGame();
    return_to_menu = true;
  }
}

void displayMenu() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(WHITE);
  int centerX = 48;
  int startY = 16;
  
  display.setCursor(centerX - 32, startY);
  display.print("1. Start ");
  display.setCursor(centerX - 32, startY + 8);
  display.print("2. Paddle Size ");
  display.setCursor(centerX + 56, startY + 8);
  display.print(paddle_width);
  display.setCursor(centerX - 32, startY + 16);
  display.print("3. Ball Speed ");
  display.setCursor(centerX + 50, startY + 16);
  display.print(selected_ball_speed);
  display.setCursor(centerX - 40, startY + 8 * (menu_selection - 1));
  display.print(">");
  display.display();
}

void handleMenu() {
  int joystickY = analogRead(JOYSTICK_Y);
  float normalizedY = map(joystickY, 0, 1023, -1, 1);

  menu_selection -= int(normalizedY);
  if (menu_selection < 1) {
    menu_selection = 3;
  }
  if (menu_selection > 3) {
    menu_selection = 1;
  }

  int joystickX = analogRead(JOYSTICK_X);
  float normalizedX = map(joystickX, 0, 1023, -1, 1);

  if (menu_selection == 1) {
    if (abs(normalizedX) > 0.1) {
      game_started = true;
      resetGame();
      return_to_menu = false;
    }
  } else if (menu_selection == 2) {
    handlePaddleSizeMenu();
  } else if (menu_selection == 3) {
    handleBallSpeedMenu();
  }
}

void handlePaddleSizeMenu() {
  int joystickX = analogRead(JOYSTICK_X);
  float normalizedX = map(joystickX, 0, 1023, -1, 1);
  paddle_width -= int(normalizedX * 8);

  if (paddle_width < PADDLE_WIDTH_MIN) {
    paddle_width = PADDLE_WIDTH_MIN;
  }
  if (paddle_width > PADDLE_WIDTH_MAX) {
    paddle_width = PADDLE_WIDTH_MAX;
  }

  display.display();
}

void handleBallSpeedMenu() {
  int joystickX = analogRead(JOYSTICK_X);
  float normalizedX = map(joystickX, 0, 1023, -1, 1);
  selected_ball_speed -= normalizedX * 1;

  if (selected_ball_speed < 1) {
    selected_ball_speed = 1;
  }
  if (selected_ball_speed > 4.0) {
    selected_ball_speed = 4.0;
  }

  display.display();
}

void handleGameplay() {
  int joystickX = analogRead(JOYSTICK_X);
  float normalizedX = map(joystickX, 0, 1023, -1, 1);
  paddle_x -= int(normalizedX * paddle_speed);

  if (paddle_x < 0) {
    paddle_x = 0;
  }
  if (paddle_x > 128 - paddle_width) {
    paddle_x = 128 - paddle_width;
  }

  int joystickY = analogRead(JOYSTICK_Y);
  float normalizedY = map(joystickY, 0, 1023, -1, 1);

  if (abs(normalizedY) > 0.1) {
    game_started = false;
    return_to_menu = true;
  }

  moveBall();
}

void moveBall() {
  ball_x += ball_speed_x;
  ball_y += ball_speed_y;

  if (ball_x < 0 || ball_x > 128.0 - BALL_SIZE) {
    ball_speed_x = -ball_speed_x;
  }
  if (ball_y < 0) {
    ball_speed_y = -ball_speed_y;
  }
  if (ball_y > 64.0 - BALL_SIZE) {
    game_over = true;
  }

  if (ball_x + BALL_SIZE >= paddle_x && ball_x <= paddle_x + paddle_width && ball_y + BALL_SIZE >= paddle_y) {
    float relativeIntersectX = (paddle_x + paddle_width / 2) - ball_x;
    angle = (PI / 4) * (relativeIntersectX / (paddle_width / 2));
    ball_speed_y = -ball_speed_y;
    ball_y = paddle_y - BALL_SIZE;
    float magnitude = sqrt(ball_speed_x * ball_speed_x + ball_speed_y * ball_speed_y);
    ball_speed_x = sin(angle) * magnitude;
    ball_speed_y = -cos(angle) * magnitude;
  }

  int ball_grid_x = ball_x / (BRICK_WIDTH + BRICK_GAP);
  int ball_grid_y = ball_y / (BRICK_HEIGHT + BRICK_GAP);

  if (ball_grid_x >= 0 && ball_grid_x < BRICK_COLUMNS && ball_grid_y >= 0 && ball_grid_y < BRICK_ROWS) {
    if (bricks[ball_grid_x][ball_grid_y]) {
      ball_speed_y = -ball_speed_y;
      bricks[ball_grid_x][ball_grid_y] = false;
      brick_colors[ball_grid_x][ball_grid_y] = BLACK;
      bool all_bricks_cleared = true;
      for (int i = 0; i < BRICK_COLUMNS; i++) {
        for (int j = 0; j < BRICK_ROWS; j++) {
          if (bricks[i][j]) {
            all_bricks_cleared = false;
            break;
          }
        }
        if (!all_bricks_cleared) break;
      }
      if (all_bricks_cleared) {
        win = true;
        game_over = true;
      }
    }
  }
}

void drawGame() {
  display.clearDisplay();
  if (game_over && game_started) {
    display.setTextSize(2);
    display.setTextColor(WHITE);
    display.setCursor(16, 24);
    display.fillRect(0, 0, 128, 64, BLACK);
    if (win) {
      display.print("You win!");
    } else {
      display.print("You lose!");
    }
  } else {
    for (int i = 0; i < BRICK_COLUMNS; i++) {
      for (int j = 0; j < BRICK_ROWS; j++) {
        if (bricks[i][j]) {
          int brick_x = i * (BRICK_WIDTH + BRICK_GAP);
          int brick_y = j * (BRICK_HEIGHT + BRICK_GAP);
          display.fillRect(brick_x, brick_y, BRICK_WIDTH, BRICK_HEIGHT, brick_colors[i][j]);
        }
      }
    }
    display.fillRect(paddle_x, paddle_y, paddle_width, PADDLE_HEIGHT, WHITE);
    display.fillCircle(ball_x, ball_y, BALL_SIZE, WHITE);
  }

  display.display();
  delay(10);
}

void resetGame() {
  game_over = false;
  win = false;
  paddle_x = (128 - paddle_width) / 2;
  paddle_y = 60;
  ball_x = 64.0;
  ball_y = 32.0;
  ball_speed_x = selected_ball_speed;
  ball_speed_y = -selected_ball_speed;

  for (int i = 0; i < BRICK_COLUMNS; i++) {
    for (int j = 0; j < BRICK_ROWS; j++) {
      bricks[i][j] = true;
      brick_colors[i][j] = WHITE;
    }
  }
}