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

// OLED SSD1306 128x64 — plansza gry
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64

Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, -1);

// LCD1602 I2C — punkty i stan gry
LiquidCrystal_I2C lcd(0x27, 16, 2);

// Przyciski
const int BTN_UP = 2;
const int BTN_DOWN = 3;
const int BTN_LEFT = 4;
const int BTN_RIGHT = 5;

// Plansza gry na OLED
const int CELL_SIZE = 8;
const int GRID_W = SCREEN_WIDTH / CELL_SIZE;   // 16
const int GRID_H = SCREEN_HEIGHT / CELL_SIZE;  // 8

const int MAX_SNAKE = 80;
const unsigned long MOVE_DELAY = 300;

struct Cell {
  int x;
  int y;
};

Cell snake[MAX_SNAKE];
int snakeLength = 3;

Cell fruit;

int score = 0;

enum Direction {
  DIR_UP,
  DIR_DOWN,
  DIR_LEFT,
  DIR_RIGHT
};

Direction direction = DIR_RIGHT;
Direction nextDirection = DIR_RIGHT;

unsigned long lastMoveTime = 0;
bool gameOver = false;
bool waitingForStart = true;

bool sameCell(Cell a, Cell b) {
  return a.x == b.x && a.y == b.y;
}

bool cellOnSnake(Cell c) {
  for (int i = 0; i < snakeLength; i++) {
    if (sameCell(c, snake[i])) {
      return true;
    }
  }
  return false;
}

bool anyButtonPressed() {
  return digitalRead(BTN_UP) == LOW ||
         digitalRead(BTN_DOWN) == LOW ||
         digitalRead(BTN_LEFT) == LOW ||
         digitalRead(BTN_RIGHT) == LOW;
}

void spawnFruit() {
  do {
    fruit.x = random(0, GRID_W);
    fruit.y = random(0, GRID_H);
  } while (cellOnSnake(fruit));
}

void lcdShowStart() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("SNAKE");
  lcd.setCursor(0, 1);
  lcd.print("Press button");
}

void lcdShowPlaying() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("Status: PLAY");

  lcd.setCursor(0, 1);
  lcd.print("Points: ");
  lcd.print(score);
}

void lcdShowGameOver() {
  lcd.clear();
  lcd.setCursor(0, 0);
  lcd.print("GAME OVER");

  lcd.setCursor(0, 1);
  lcd.print("Points: ");
  lcd.print(score);
}

void resetGame() {
  snakeLength = 3;
  score = 0;

  snake[0] = { GRID_W / 2, GRID_H / 2 };
  snake[1] = { GRID_W / 2 - 1, GRID_H / 2 };
  snake[2] = { GRID_W / 2 - 2, GRID_H / 2 };

  direction = DIR_RIGHT;
  nextDirection = DIR_RIGHT;

  gameOver = false;
  lastMoveTime = millis();

  spawnFruit();
}

void readButtons() {
  if (digitalRead(BTN_UP) == LOW && direction != DIR_DOWN) {
    nextDirection = DIR_UP;
  }

  if (digitalRead(BTN_DOWN) == LOW && direction != DIR_UP) {
    nextDirection = DIR_DOWN;
  }

  if (digitalRead(BTN_LEFT) == LOW && direction != DIR_RIGHT) {
    nextDirection = DIR_LEFT;
  }

  if (digitalRead(BTN_RIGHT) == LOW && direction != DIR_LEFT) {
    nextDirection = DIR_RIGHT;
  }
}

void moveSnake() {
  direction = nextDirection;

  Cell newHead = snake[0];

  if (direction == DIR_UP) {
    newHead.y--;
  } 
  else if (direction == DIR_DOWN) {
    newHead.y++;
  } 
  else if (direction == DIR_LEFT) {
    newHead.x--;
  } 
  else if (direction == DIR_RIGHT) {
    newHead.x++;
  }

  // Kolizja ze ścianą
  if (newHead.x < 0 || newHead.x >= GRID_W || newHead.y < 0 || newHead.y >= GRID_H) {
    gameOver = true;
    lcdShowGameOver();
    return;
  }

  bool ateFruit = sameCell(newHead, fruit);

  // Kolizja z samym sobą
  int collisionLimit = ateFruit ? snakeLength : snakeLength - 1;

  for (int i = 0; i < collisionLimit; i++) {
    if (sameCell(newHead, snake[i])) {
      gameOver = true;
      lcdShowGameOver();
      return;
    }
  }

  if (ateFruit && snakeLength < MAX_SNAKE) {
    snakeLength++;
    score++;
    lcdShowPlaying();
  }

  // Przesunięcie ciała węża
  for (int i = snakeLength - 1; i > 0; i--) {
    snake[i] = snake[i - 1];
  }

  snake[0] = newHead;

  if (ateFruit) {
    spawnFruit();
  }
}

void drawBoard() {
  display.clearDisplay();

  // Ramka planszy
  display.drawRect(0, 0, SCREEN_WIDTH, SCREEN_HEIGHT, SSD1306_WHITE);

  // Owoc — pusty kwadrat
  display.drawRect(
    fruit.x * CELL_SIZE + 1,
    fruit.y * CELL_SIZE + 1,
    CELL_SIZE - 2,
    CELL_SIZE - 2,
    SSD1306_WHITE
  );

  // Wąż — pełne kwadraty
  for (int i = 0; i < snakeLength; i++) {
    display.fillRect(
      snake[i].x * CELL_SIZE + 1,
      snake[i].y * CELL_SIZE + 1,
      CELL_SIZE - 2,
      CELL_SIZE - 2,
      SSD1306_WHITE
    );
  }

  display.display();
}

void drawStartBoard() {
  display.clearDisplay();

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  display.setCursor(34, 20);
  display.println("SNAKE");

  display.setCursor(10, 40);
  display.println("Press any button");

  display.display();
}

void drawGameOverBoard() {
  display.clearDisplay();

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);

  display.setCursor(36, 20);
  display.println("GAME OVER");

  display.setCursor(14, 40);
  display.println("Press to restart");

  display.display();
}

void setup() {
  pinMode(BTN_UP, INPUT_PULLUP);
  pinMode(BTN_DOWN, INPUT_PULLUP);
  pinMode(BTN_LEFT, INPUT_PULLUP);
  pinMode(BTN_RIGHT, INPUT_PULLUP);

  // Start I2C
  Wire.begin();

  // OLED
  if (!display.begin(SSD1306_SWITCHCAPVCC, 0x3C)) {
    while (true);
  }

  // LCD1602
  lcd.init();
  lcd.backlight();

  randomSeed(analogRead(A0));

  resetGame();

  waitingForStart = true;
  lcdShowStart();
  drawStartBoard();
}

void loop() {
  // Ekran startowy
  if (waitingForStart) {
    if (anyButtonPressed()) {
      delay(250);
      waitingForStart = false;
      resetGame();
      lcdShowPlaying();
      drawBoard();
    }

    return;
  }

  // Game over
  if (gameOver) {
    drawGameOverBoard();

    if (anyButtonPressed()) {
      delay(250);
      resetGame();
      waitingForStart = false;
      lcdShowPlaying();
      drawBoard();
    }

    return;
  }

  // Gra
  readButtons();

  if (millis() - lastMoveTime >= MOVE_DELAY) {
    lastMoveTime = millis();
    moveSnake();

    if (!gameOver) {
      drawBoard();
    }
  }
}