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

// OLED setup
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define OLED_RESET -1
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// Button pin variables
int buttonUp = D10;     // Replace with your GPIO pin number
int buttonDown = D9;   // Replace with your GPIO pin number
int buttonLeft = D8;   // Replace with your GPIO pin number
int buttonRight = D7;  // Replace with your GPIO pin number

// Snake game variables
#define GRID_SIZE 4     // Size of one grid square
#define GRID_WIDTH (SCREEN_WIDTH / GRID_SIZE)
#define GRID_HEIGHT (SCREEN_HEIGHT / GRID_SIZE)

int snakeX[128], snakeY[64]; // Max size of the snake
int snakeLength = 3;         // Initial length of the snake
int foodX, foodY;            // Food position
int direction = 0;           // 0 = Up, 1 = Right, 2 = Down, 3 = Left
bool gameOver = false;

void setup() {
  // Initialize the display
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    while (true); // Loop forever if OLED initialization fails
  }
  display.clearDisplay();

  // Initialize buttons
  pinMode(buttonUp, INPUT_PULLUP);
  pinMode(buttonDown, INPUT_PULLUP);
  pinMode(buttonLeft, INPUT_PULLUP);
  pinMode(buttonRight, INPUT_PULLUP);

  // Initialize snake position
  for (int i = 0; i < snakeLength; i++) {
    snakeX[i] = GRID_WIDTH / 2;
    snakeY[i] = (GRID_HEIGHT / 2) + i;
  }

  // Place initial food
  placeFood();
}

void loop() {
  if (!gameOver) {
    handleInput();
    moveSnake();
    checkCollision();
    drawGame();
    delay(200); // Adjust speed of the game
  } else {
    displayGameOver();
  }
}

// Handle button presses
void handleInput() {
  if (digitalRead(buttonUp) == LOW && direction != 2) direction = 0;
  if (digitalRead(buttonRight) == LOW && direction != 3) direction = 1;
  if (digitalRead(buttonDown) == LOW && direction != 0) direction = 2;
  if (digitalRead(buttonLeft) == LOW && direction != 1) direction = 3;
}

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

  // Move head
  if (direction == 0) snakeY[0]--; // Up
  if (direction == 1) snakeX[0]++; // Right
  if (direction == 2) snakeY[0]++; // Down
  if (direction == 3) snakeX[0]--; // Left

  // Check if the snake eats the food
  if (snakeX[0] == foodX && snakeY[0] == foodY) {
    snakeLength++;
    placeFood();
  }
}

// Place food at a random location
void placeFood() {
  foodX = random(0, GRID_WIDTH);
  foodY = random(0, GRID_HEIGHT);
}

// Check for collisions
void checkCollision() {
  // Check wall collision
  if (snakeX[0] < 0 || snakeX[0] >= GRID_WIDTH || snakeY[0] < 0 || snakeY[0] >= GRID_HEIGHT) {
    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 screen
void drawGame() {
  display.clearDisplay();

  // Draw food
  display.fillRect(foodX * GRID_SIZE, foodY * GRID_SIZE, GRID_SIZE, GRID_SIZE, SSD1306_WHITE);

  // Draw snake
  for (int i = 0; i < snakeLength; i++) {
    display.fillRect(snakeX[i] * GRID_SIZE, snakeY[i] * GRID_SIZE, GRID_SIZE, GRID_SIZE, SSD1306_WHITE);
  }

  display.display();
}

// Display Game Over message
void displayGameOver() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(10, 20);
  display.println("GAME OVER");
  display.display();
  delay(5000); // Show game over screen for 5 seconds
  ESP.restart(); // Restart the game
}