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

// ---------- OLED Setup ----------
#define SCREEN_WIDTH 128
#define SCREEN_HEIGHT 64
#define SCREEN_ADDRESS 0x3C

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

// ---------- Snake Setup ----------
#define SNAKE_SIZE 3   // Increased size for easier overlap
#define MAX_LENGTH 100

int snakeX[MAX_LENGTH];
int snakeY[MAX_LENGTH];
int snakeLength = 5;
int dirX = 1;
int dirY = 0;

int foodX, foodY;

// ---------- Touch Inputs ----------
#define N_TOUCH 4
#define THRESHOLD 30

// Pin order: UP, DOWN, LEFT, RIGHT
int touch_pins[N_TOUCH] = {26, 2, 27, 1};
int touch_values[N_TOUCH] = {0, 0, 0, 0};
bool pin_touched_now[N_TOUCH] = {false,false,false,false};
bool pin_touched_past[N_TOUCH] = {false,false,false,false};

// ---------- Functions ----------
void spawnFood() {
  foodX = (random(SCREEN_WIDTH / SNAKE_SIZE)) * SNAKE_SIZE;
  foodY = (random(SCREEN_HEIGHT / SNAKE_SIZE)) * SNAKE_SIZE;
}

void update_touch() {
  int t;
  int t_max = 200;
  int p;
  for (int i = 0; i < N_TOUCH; i++) {
    p = touch_pins[i];
    pinMode(p, OUTPUT);
    digitalWrite(p, LOW);
    delayMicroseconds(25);
    noInterrupts();
    pinMode(p, INPUT_PULLUP);
    t = 0;
    while (!digitalRead(p) && t < t_max) t++;
    interrupts();
    touch_values[i] = t;
    pin_touched_past[i] = pin_touched_now[i];
    pin_touched_now[i] = touch_values[i] > THRESHOLD;
  }
}

void gameOver() {
  display.clearDisplay();
  display.setTextSize(2);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(15, 20);
  display.println("Game Over");
  display.display();
  delay(2000);

  // Reset snake
  snakeLength = 5;
  snakeX[0] = SCREEN_WIDTH / 2;
  snakeY[0] = SCREEN_HEIGHT / 2;
  for (int i = 1; i < snakeLength; i++) {
    snakeX[i] = snakeX[0] - i * SNAKE_SIZE;
    snakeY[i] = snakeY[0];
  }
  dirX = 1;
  dirY = 0;
  spawnFood();
}

// ---------- Setup ----------
void setup() {
  Serial.begin(115200);

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

  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(20, 20);
  display.println("Snake starting...");
  display.display();
  delay(2000);

  randomSeed(analogRead(0));

  // Initialize snake in the middle
  snakeX[0] = SCREEN_WIDTH / 2;
  snakeY[0] = SCREEN_HEIGHT / 2;
  for (int i = 1; i < snakeLength; i++) {
    snakeX[i] = snakeX[0] - i * SNAKE_SIZE;
    snakeY[i] = snakeY[0];
  }

  spawnFood();
}

// ---------- Main Loop ----------
void loop() {
  update_touch();

  // Directions: UP=0, DOWN=1, LEFT=2, RIGHT=3
  if (pin_touched_now[0] && !pin_touched_past[0] && dirY == 0) { dirX = 0; dirY = -1; }
  if (pin_touched_now[1] && !pin_touched_past[1] && dirY == 0) { dirX = 0; dirY = 1; }
  if (pin_touched_now[2] && !pin_touched_past[2] && dirX == 0) { dirX = -1; dirY = 0; }
  if (pin_touched_now[3] && !pin_touched_past[3] && dirX == 0) { dirX = 1; dirY = 0; }

  // Move snake body
  for (int i = snakeLength - 1; i > 0; i--) {
    snakeX[i] = snakeX[i - 1];
    snakeY[i] = snakeY[i - 1];
  }
  snakeX[0] += dirX * SNAKE_SIZE;
  snakeY[0] += dirY * SNAKE_SIZE;

  // Check collision with food (overlap)
  if (snakeX[0] < foodX + SNAKE_SIZE &&
      snakeX[0] + SNAKE_SIZE > foodX &&
      snakeY[0] < foodY + SNAKE_SIZE &&
      snakeY[0] + SNAKE_SIZE > foodY) {

    // Extend snake
    if (snakeLength < MAX_LENGTH) {
      snakeX[snakeLength] = snakeX[snakeLength - 1];
      snakeY[snakeLength] = snakeY[snakeLength - 1];
      snakeLength++;
    }
    spawnFood();
  }

  // Wall collision
  if (snakeX[0] < 0 || snakeX[0] >= SCREEN_WIDTH || snakeY[0] < 0 || snakeY[0] >= SCREEN_HEIGHT) {
    gameOver();
  }

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

  // Draw everything
  display.clearDisplay();
  display.fillRect(foodX, foodY, SNAKE_SIZE, SNAKE_SIZE, SSD1306_WHITE);
  for (int i = 0; i < snakeLength; i++) {
    display.fillRect(snakeX[i], snakeY[i], SNAKE_SIZE, SNAKE_SIZE, SSD1306_WHITE);
  }
  display.display();

  delay(150);
}