#include <U8g2lib.h>

// Pin definitions for buttons
#define UP_BUTTON 19
#define DOWN_BUTTON 18
#define LEFT_BUTTON 17
#define RIGHT_BUTTON 16
#define RESET_BUTTON 25 // Pin for the reset button

// Display setup
U8G2_SSD1306_128X64_NONAME_F_HW_I2C u8g2(U8G2_R0, /* reset=*/U8X8_PIN_NONE);

// Game variables
int snakeX[100] = {64}; // X coordinates of the snake body
int snakeY[100] = {32}; // Y coordinates of the snake body
int snakeLength = 5;    // Initial length of the snake
int foodX, foodY;       // Food position
int dirX = 1, dirY = 0; // Snake direction (initially moving right)
bool gameOver = false;

// Function prototypes
void generateFood();
void updateSnake();
void checkCollision();
void drawGame();
void readButtons();
void resetGame();
void displayGameOver();
bool isButtonPressed(int pin);

void setup() {
  // Initialize display
  u8g2.begin();
  u8g2.setFont(u8g2_font_6x10_tr);

  // Initialize buttons as input with pulldown
  pinMode(UP_BUTTON, INPUT_PULLDOWN);
  pinMode(DOWN_BUTTON, INPUT_PULLDOWN);
  pinMode(LEFT_BUTTON, INPUT_PULLDOWN);
  pinMode(RIGHT_BUTTON, INPUT_PULLDOWN);
  pinMode(RESET_BUTTON, INPUT_PULLDOWN);  // Reset button setup

  // Generate initial food position
  generateFood();
}

void loop() {
  if (isButtonPressed(RESET_BUTTON)) {
    resetGame();  // Reset game state when restart button is pressed
  }

  if (gameOver) {
    // Display "Game Over" screen and the final score
    displayGameOver();
    return;
  }

  readButtons();  // Read button input to update direction
  updateSnake();  // Move the snake
  checkCollision(); // Check for collisions
  drawGame();     // Render the game
  delay(150);     // Control snake speed
}

// Generate food at a random position
void generateFood() {
  foodX = (random(0, 128 / 4)) * 4; // Align to grid (4x4 blocks)
  foodY = (random(0, 64 / 4)) * 4;
}

// Update snake position
void updateSnake() {
  // Move the body of the snake
  for (int i = snakeLength - 1; i > 0; i--) {
    snakeX[i] = snakeX[i - 1];
    snakeY[i] = snakeY[i - 1];
  }

  // Move the head of the snake
  snakeX[0] += dirX * 4;
  snakeY[0] += dirY * 4;

  // Check if the snake eats the food
  if (snakeX[0] == foodX && snakeY[0] == foodY) {
    snakeLength++; // Grow the snake
    generateFood(); // Generate a new food
  }
}

// Check for collisions with walls or itself
void checkCollision() {
  // Check wall collision
  if (snakeX[0] < 0 || snakeX[0] >= 128 || snakeY[0] < 0 || snakeY[0] >= 64) {
    gameOver = true;
  }

  // Check self-collision
  for (int i = 1; i < snakeLength; i++) {
    if (snakeX[0] == snakeX[i] && snakeY[0] == snakeY[i]) {
      gameOver = true;
    }
  }
}

// Draw the game on the display
void drawGame() {
  u8g2.clearBuffer();

  // Draw the snake
  for (int i = 0; i < snakeLength; i++) {
    u8g2.drawBox(snakeX[i], snakeY[i], 4, 4); // Snake segments are 4x4 blocks
  }

  // Draw the food
  u8g2.drawBox(foodX, foodY, 4, 4);

  u8g2.sendBuffer();
}

// Read button input to change direction
void readButtons() {
  if (digitalRead(UP_BUTTON) == HIGH && dirY == 0) {
    dirX = 0; dirY = -1; // Move up
  } else if (digitalRead(DOWN_BUTTON) == HIGH && dirY == 0) {
    dirX = 0; dirY = 1; // Move down
  } else if (digitalRead(LEFT_BUTTON) == HIGH && dirX == 0) {
    dirX = -1; dirY = 0; // Move left
  } else if (digitalRead(RIGHT_BUTTON) == HIGH && dirX == 0) {
    dirX = 1; dirY = 0; // Move right
  }
}

// Reset the game (clear snake, reset length, generate new food)
void resetGame() {
  snakeLength = 5;
  dirX = 1; dirY = 0; // Reset direction to right
  gameOver = false;

  // Clear the snake's position
  for (int i = 0; i < 100; i++) {
    snakeX[i] = 64;
    snakeY[i] = 32;
  }

  // Generate new food position
  generateFood();
}

// Display the Game Over screen and show the score
void displayGameOver() {
  u8g2.clearBuffer();

  // Display "Game Over" message
  u8g2.drawStr(30, 30, "Game Over!");

  // Display the score (snake length - 5, since length starts at 5)
  String score = "Score: " + String(snakeLength - 5);
  u8g2.drawStr(20, 45, score.c_str());

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

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

  resetGame();
}

// Check if the button is pressed and debounce it
bool isButtonPressed(int pin) {
  static unsigned long lastPressTime = 0;
  unsigned long currentMillis = millis();

  // Check if button is pressed
  if (digitalRead(pin) == HIGH && currentMillis - lastPressTime > 1000) { // Debounce time of 500ms
    lastPressTime = currentMillis;
    return true;
  }

  return false;
}