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

#define SCREEN_WIDTH        128 // OLED display width, in pixels
#define SCREEN_HEIGHT        64 // OLED display height, in pixels
#define OLED_RESET            4 // Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

const byte buttonPins[] = {4, 2, 5, 3}; // LEFT, UP, RIGHT, DOWN

typedef enum { LEFT, UP, RIGHT, DOWN } Direction;
Direction dir = RIGHT, newDir = RIGHT;

#define SNAKE_PIECE_SIZE     3
#define MAX_SANKE_LENGTH   165
#define MAP_SIZE_X          43
#define MAP_SIZE_Y          22
#define STARTING_SNAKE_SIZE  5

int8_t snake[MAX_SANKE_LENGTH][2];
uint8_t snake_length;
int8_t fruit[2];
bool gameOver = false;

void setup() {
  Serial.begin(9600);

  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3D)) {
    Serial.println(F("SSD1306 allocation failed"));
    for (;;);
  }

  for (byte i = 0; i < 4; i++) pinMode(buttonPins[i], INPUT_PULLUP);
  randomSeed(analogRead(A0));
  resetGame();
}

void resetGame() {
  snake_length = STARTING_SNAKE_SIZE;
  for (int i = 0; i < snake_length; i++) {
    snake[i][0] = MAP_SIZE_X / 2 - i;
    snake[i][1] = MAP_SIZE_Y / 2;
  }
  dir = newDir = RIGHT;
  fruit[0] = random(0, MAP_SIZE_X);
  fruit[1] = random(0, MAP_SIZE_Y);
  gameOver = false;

  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 20);
  display.println(F("Press any\nbutton to\nstart!"));
  display.display();
}

void loop() {
  static bool running = false;

  // Handle game over state
  if (gameOver) {
    for (byte i = 0; i < 4; i++) {
      if (digitalRead(buttonPins[i]) == LOW) {
        delay(500); // debounce delay
        resetGame();
        return;
      }
    }
    return; // Wait for button press to restart
  }

  // Start game on button press
  if (!running) {
    for (byte i = 0; i < 4; i++)
      if (digitalRead(buttonPins[i]) == LOW) running = true;
    return;
  }

  // Read direction
  for (byte i = 0; i < 4; i++) {
    if (digitalRead(buttonPins[i]) == LOW && i != ((int)dir + 2) % 4) {
      newDir = (Direction)i;
      break;
    }
  }

  // Move snake
  delay(100);
  dir = newDir;
  int8_t x = snake[0][0], y = snake[0][1];
  if (dir == LEFT) x--; else if (dir == RIGHT) x++;
  else if (dir == UP) y--; else if (dir == DOWN) y++;

  // Collision check
  for (int i = 1; i < snake_length; i++)
    if (x == snake[i][0] && y == snake[i][1]) {
      gameOver = true;
      return drawGameOver();
    }
  if (x < 0 || y < 0 || x >= MAP_SIZE_X || y >= MAP_SIZE_Y) {
    gameOver = true;
    return drawGameOver();
  }

  // Move snake
  for (int i = snake_length - 1; i > 0; i--) {
    snake[i][0] = snake[i - 1][0];
    snake[i][1] = snake[i - 1][1];
  }
  snake[0][0] = x;
  snake[0][1] = y;

  // Eat fruit
  if (x == fruit[0] && y == fruit[1]) {
    if (snake_length < MAX_SANKE_LENGTH) snake_length++;
    fruit[0] = random(0, MAP_SIZE_X);
    fruit[1] = random(0, MAP_SIZE_Y);
  }

  // Draw game
  display.clearDisplay();
  display.drawRect(fruit[0] * SNAKE_PIECE_SIZE, fruit[1] * SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  for (int i = 0; i < snake_length; i++) {
    display.fillRect(snake[i][0] * SNAKE_PIECE_SIZE, snake[i][1] * SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  }
  display.setCursor(2, 2);
  display.print(F("Score:"));
  display.print(snake_length - STARTING_SNAKE_SIZE);
  display.display();
}

void drawGameOver() {
  display.clearDisplay();
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 10);
  display.println(F("GAME OVER"));

  display.setCursor(2, 30);
  display.print(F("Score: "));
  display.println(snake_length - STARTING_SNAKE_SIZE);

  display.setCursor(2, 50);
  display.println(F("Press any button\nto restart!"));

  display.display();
}