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

// OLED display settings
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

// Define OLED display object
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// Define buttons and snake block size
const int LEFT_BUTTON = 5;
const int UP_BUTTON = 3;
const int DOWN_BUTTON = 4;
const int RIGHT_BUTTON = 2;

//***************
// SNEK
//***************
bool playsnek = 0;

// Define game variables
int deaths = 0;
int score = 0;
int highscore = 0;
int snake_x = 0;
int snake_y = 0;
int snake_dir_x = 1;
int snake_dir_y = 0;
int food_x = 0;
int food_y = 0;
boolean is_food_eaten = true;
const int MAX_SEGMENTS = 50; // 100 segments might crash memory
int segment_x[MAX_SEGMENTS];
int segment_y[MAX_SEGMENTS];
int num_segments = 3; // Starting with head and 2 additional segments
float speed_multiplier = 1.0;
const int SNAKE_BLOCK_SIZE = 4;
const float SPEED_INCREMENT = 0.2;
const int BASE_DELAY_TIME = 100;
const int SCREEN_LEFT_EDGE = 0;
const int SCREEN_RIGHT_EDGE = SCREEN_WIDTH - SNAKE_BLOCK_SIZE;
const int SCREEN_TOP_EDGE = 0;
const int SCREEN_BOTTOM_EDGE = SCREEN_HEIGHT - SNAKE_BLOCK_SIZE;

//***************
// PING
//***************
bool playping = 0;

// Paddle dimensions
const int PADDLE_WIDTH = 4;
const int PADDLE_HEIGHT = 20;

// Define paddle positions
int player1_paddle_y = SCREEN_HEIGHT / 2;
int player2_paddle_y = SCREEN_HEIGHT / 2;

// Ball dimensions and speed
const int BALL_SIZE = 4;
const int BALL_SPEED = 2;

// Game variables
int player1_score = 0;
int player2_score = 0;

// Ball position and direction
int ball_x = SCREEN_WIDTH / 2;
int ball_y = SCREEN_HEIGHT / 2;
int ball_dx = BALL_SPEED;
int ball_dy = BALL_SPEED;

void setup() {
  // OLED display
  display.begin(SSD1306_SWITCHCAPVCC, 0x3C);
  display.clearDisplay();
  display.display();

  // Buttons
  pinMode(LEFT_BUTTON, INPUT_PULLUP);
  pinMode(UP_BUTTON, INPUT_PULLUP);
  pinMode(DOWN_BUTTON, INPUT_PULLUP);
  pinMode(RIGHT_BUTTON, INPUT_PULLUP);

  // Select
  display.setCursor(0, 10);
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.println("PRESS LEFT FOR SNEK");
  display.print("\nPRESS RIGHT FOR PING");
  display.display();

  while (digitalRead(LEFT_BUTTON) == HIGH && digitalRead(RIGHT_BUTTON) == HIGH) {
    if (!digitalRead(RIGHT_BUTTON))
      playping = 1;
    if (!digitalRead(LEFT_BUTTON))
      playsnek = 1;
  }
  delay(50); // button debounce
  display.clearDisplay();

  //********************
  // SNEK setup
  //********************
  if (playsnek) {

    // Print initial screen with title and prompt
    title("SNEK");

    // Set initial snake position and add to segment arrays
    resetSnake();

    // Draw snake graphic on welcome screen
    display.fillRect(100, 8, 8, 4, WHITE);
    display.fillRect(96, 12, 8, 4, WHITE);
    display.fillRect(92, 16, 8, 4, WHITE);
    display.fillRect(88, 20, 8, 4, WHITE);
    display.display();

    // Wait for any button to be pressed
    pressanybutton();
  } // SNEK setup END

  //********************
  // PING setup
  //********************
  if (playping) {
    // Print initial screen with title and prompt
    title("PING");

    // Wait for any button to be pressed
    pressanybutton();

    // Draw initial screen with paddles and ball
    display.clearDisplay();
    display.fillRect(0, player1_paddle_y - PADDLE_HEIGHT / 2, PADDLE_WIDTH, PADDLE_HEIGHT, WHITE);
    display.fillRect(SCREEN_WIDTH - PADDLE_WIDTH, player2_paddle_y - PADDLE_HEIGHT / 2, PADDLE_WIDTH, PADDLE_HEIGHT, WHITE);
    display.fillRect(ball_x - BALL_SIZE / 2, ball_y - BALL_SIZE / 2, BALL_SIZE, BALL_SIZE, WHITE);
    update_scoreboard();
    display.display();
  } // PING setup END
}

void loop() {

  //********************
  // SNEK
  //********************
  if (playsnek) {

    boolean is_collision = false;

    // Clear display before drawing new elements
    display.clearDisplay();

    // Calculate delay time based on current speed multiplier
    int delay_time = (int)(BASE_DELAY_TIME / speed_multiplier);

    // Check button input and update snake direction
    if (digitalRead(LEFT_BUTTON) == LOW) {
      snake_dir_x = -1;
      snake_dir_y = 0;
    } else if (digitalRead(UP_BUTTON) == LOW) {
      snake_dir_x = 0;
      snake_dir_y = -1;
    } else if (digitalRead(DOWN_BUTTON) == LOW) {
      snake_dir_x = 0;
      snake_dir_y = 1;
    } else if (digitalRead(RIGHT_BUTTON) == LOW) {
      snake_dir_x = 1;
      snake_dir_y = 0;
    }

    // Erase previous snake head position
    display.fillRect(segment_x[num_segments - 1], segment_y[num_segments - 1], SNAKE_BLOCK_SIZE, SNAKE_BLOCK_SIZE, BLACK);

    // Move snake segments
    for (int i = num_segments - 1; i > 0; i--) {
      segment_x[i] = segment_x[i - 1];
      segment_y[i] = segment_y[i - 1];
    }

    // Update snake head position
    snake_x += snake_dir_x * SNAKE_BLOCK_SIZE;
    snake_y += snake_dir_y * SNAKE_BLOCK_SIZE;

    // Wrap around to opposite edge if going out of bounds
    if (snake_x < SCREEN_LEFT_EDGE) {
      snake_x = SCREEN_RIGHT_EDGE;
    } else if (snake_x > SCREEN_RIGHT_EDGE) {
      snake_x = SCREEN_LEFT_EDGE;
    }
    if (snake_y < SCREEN_TOP_EDGE) {
      snake_y = SCREEN_BOTTOM_EDGE;
    } else if (snake_y > SCREEN_BOTTOM_EDGE) {
      snake_y = SCREEN_TOP_EDGE;
    }

    // Check for food collision
    if (snake_x == food_x && snake_y == food_y) {
      // Increase speed and reset food position
      speed_multiplier += SPEED_INCREMENT;
      is_food_eaten = true;
      tone(48, 40, 30);
      score++;

      // Update high score if necessary
      if (score > highscore) {
        highscore = score;
      }

      // Add new segment to snake if possible
      if (num_segments < MAX_SEGMENTS) {
        segment_x[num_segments] = segment_x[num_segments - 1];
        segment_y[num_segments] = segment_y[num_segments - 1];
        num_segments++;
      }
    }

    // Check for self-collision
    for (int i = 1; i < num_segments; i++) {
      if (snake_x == segment_x[i] && snake_y == segment_y[i]) {
        is_collision = true;
        break;
      }
    }

    // Update snake head segment
    segment_x[0] = snake_x;
    segment_y[0] = snake_y;

    // Place food if needed
    if (is_food_eaten) {
      food_x = random(SCREEN_WIDTH / SNAKE_BLOCK_SIZE) * SNAKE_BLOCK_SIZE;
      food_y = random(SCREEN_HEIGHT / SNAKE_BLOCK_SIZE) * SNAKE_BLOCK_SIZE;
      is_food_eaten = false;
    }

    // Draw snake and food
    for (int i = 0; i < num_segments; i++) {
      display.fillRect(segment_x[i], segment_y[i], SNAKE_BLOCK_SIZE, SNAKE_BLOCK_SIZE, WHITE);
    }

    display.fillCircle(food_x + SNAKE_BLOCK_SIZE / 2, food_y + SNAKE_BLOCK_SIZE / 2, SNAKE_BLOCK_SIZE / 2, WHITE);

    // Display game over message if collision detected
    if (is_collision)
      game_over();

    // Update display
    display.display();

    // Delay for smooth animation
    delay(delay_time);
  }

  //********************
  // PING
  //********************
  if (playping) {
    // Move paddles based on button inputs
    if (digitalRead(UP_BUTTON) == LOW) {
      player1_paddle_y = max(PADDLE_HEIGHT / 2, player1_paddle_y - 2);
    }
    if (digitalRead(DOWN_BUTTON) == LOW) {
      player1_paddle_y = min(SCREEN_HEIGHT - PADDLE_HEIGHT / 2, player1_paddle_y + 2);
    }

    // Move opponent paddle to follow ball
    if (ball_y < player2_paddle_y) {
      player2_paddle_y = max(PADDLE_HEIGHT / 2, player2_paddle_y - 1);
    }
    if (ball_y > player2_paddle_y) {
      player2_paddle_y = min(SCREEN_HEIGHT - PADDLE_HEIGHT / 2, player2_paddle_y + 1);
    }

    // Move ball
    ball_x += ball_dx;
    ball_y += ball_dy;

    // Check for ball collision with top or bottom walls
    if (ball_y - BALL_SIZE / 2 < 0 || ball_y + BALL_SIZE / 2 > SCREEN_HEIGHT) {
      ball_dy *= -1;
    }

    // Check for ball collision with player 1 paddle
    if (ball_x - BALL_SIZE / 2 <= PADDLE_WIDTH && abs(ball_y - player1_paddle_y) <= PADDLE_HEIGHT / 2) {
      ball_dx *= -1;
    }

    // Check for ball collision with player 2 paddle
    if (ball_x + BALL_SIZE / 2 >= SCREEN_WIDTH - PADDLE_WIDTH && abs(ball_y - player2_paddle_y) <= PADDLE_HEIGHT / 2) {
      ball_dx *= -1;
    }

    // Check for ball going out of bounds
    if (ball_x - BALL_SIZE / 2 < 0) {
      player2_score++;
      reset_ball();
    }
    if (ball_x + BALL_SIZE / 2 > SCREEN_WIDTH) {
      player1_score++;
      reset_ball();
    }

    // redraw screen with updated positions
    display.clearDisplay();
    display.fillRect(0, player1_paddle_y - PADDLE_HEIGHT / 2, PADDLE_WIDTH, PADDLE_HEIGHT, WHITE);
    display.fillRect(SCREEN_WIDTH - PADDLE_WIDTH, player2_paddle_y - PADDLE_HEIGHT / 2, PADDLE_WIDTH, PADDLE_HEIGHT, WHITE);
    display.fillRect(ball_x - BALL_SIZE / 2, ball_y - BALL_SIZE / 2, BALL_SIZE, BALL_SIZE, WHITE);
    update_scoreboard();
    display.display();
  }
}

void game_over() {
// snek
  tone(48, 400, 500);
  display.clearDisplay();
  display.setCursor(0, 10);
  display.setTextColor(WHITE);
  display.setTextSize(2);
  display.print("GAME OVER!");
  display.setTextSize(1);
  display.setCursor(0, SCREEN_HEIGHT - 30);
  display.print("Deaths: ");
  display.print(deaths + 1);
  display.setCursor(0, SCREEN_HEIGHT - 20);
  display.setTextColor(WHITE);
  display.print("Score: ");
  display.print(score);
  display.setCursor(0, SCREEN_HEIGHT - 10);
  display.print("Highscore: ");
  display.print(highscore);
  display.display();
  delay(3000);
  pressanybutton();
  reset_game();
  deaths++;
  playsnek = 0;
}

void reset_game() {
  // Reset snake position and direction
  snake_x = SCREEN_WIDTH / 2;
  snake_y = SCREEN_HEIGHT / 2;
  snake_dir_x = 0;
  snake_dir_y = -1;

  // Reset speed and score
  speed_multiplier = 1.0;
  score = 0;

  // Reset snake segments
  num_segments = 3;

  // Set initial snake position and add to segment arrays
  resetSnake();

  // Reset food position
  food_x = random(SCREEN_WIDTH / SNAKE_BLOCK_SIZE) * SNAKE_BLOCK_SIZE;
  food_y = random(SCREEN_HEIGHT / SNAKE_BLOCK_SIZE) * SNAKE_BLOCK_SIZE;
  is_food_eaten = false;
}


void resetSnake() {
  // Reset initial snake position and add to segment arrays
  segment_x[0] = snake_x;
  segment_y[0] = snake_y;
  segment_x[1] = snake_x - SNAKE_BLOCK_SIZE;
  segment_y[1] = snake_y;
  segment_x[2] = snake_x - 2 * SNAKE_BLOCK_SIZE;
  segment_y[2] = snake_y;
}

void reset_ball() {
  // ping - reset ball to center position and randomize direction
  ball_x = SCREEN_WIDTH / 2;
  ball_y = SCREEN_HEIGHT / 2;
  ball_dx = random(2) == 0 ? BALL_SPEED : -BALL_SPEED;
  ball_dy = random(2) == 0 ? BALL_SPEED : -BALL_SPEED;
}

void update_scoreboard() {
  // ping - update scoreboard with current scores
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.setCursor(SCREEN_WIDTH / 2 - 10, 0);
  display.print(player1_score);
  display.setCursor(SCREEN_WIDTH / 2 + 10, 0);
  display.print(player2_score);
  display.setCursor(SCREEN_WIDTH / 2, 0);
  display.print(":");
}

void title(char gamename[]) {
  // ping and snek
  display.setCursor(0, 10);
  display.setTextSize(2);
  display.setTextColor(WHITE);
  display.println(gamename);
  display.setCursor(0, SCREEN_HEIGHT - 25);
  display.setTextSize(1);
  display.setTextColor(WHITE);
  display.println("Press any button to");
  display.setCursor(0, SCREEN_HEIGHT - 15);
  display.println("Play");
  display.display();
}

void pressanybutton() {
  // ping and snek
  while (
    digitalRead(LEFT_BUTTON) == HIGH &&
    digitalRead(UP_BUTTON) == HIGH &&
    digitalRead(DOWN_BUTTON) == HIGH &&
    digitalRead(RIGHT_BUTTON) == HIGH
  )
    delay(50); // button debounce
}