#include <Arduino.h>
#include <U8g2lib.h>
#include <Wire.h>

U8G2_SSD1306_128X64_NONAME_1_HW_I2C u8g2(U8G2_R0, /* reset=*/ U8X8_PIN_NONE);

// Pins
#define BTN_L      D1
#define BTN_R      D3
#define BTN_START  D2

enum State { MENU, SNAKE_GAME, PONG_GAME, TIMING_GAME };
State currentState = MENU;
uint8_t menuSelection = 0;

// Input Flags
bool leftClicked = false, rightClicked = false, startClicked = false;
bool gameRunning = false;
unsigned long prevTime = 0;
int16_t score = 0;

// --- SNAKE VARS ---
#define S_GRID 4
struct Point { int8_t x, y; };
Point snake[40];
uint8_t snakeLen = 3;
Point food;
int8_t sDir = 0; // 0:R, 1:D, 2:L, 3:U
int8_t sdx[] = {1, 0, -1, 0}, sdy[] = {0, 1, 0, -1};
uint16_t snakeSpeed = 150;

// --- PONG VARS ---
float ballX, ballY, ballDX, ballDY;
int8_t paddleX = 54;

// --- CORE INPUT HANDLER ---
// This reads the buttons ONCE per loop to prevent "missing" or "double" presses.
void updateButtons() {
  static bool lastL = 1, lastR = 1, lastS = 1;
  bool cL = digitalRead(BTN_L), cR = digitalRead(BTN_R), cS = digitalRead(BTN_START);

  leftClicked = (cL == LOW && lastL == HIGH);
  rightClicked = (cR == LOW && lastR == HIGH);
  startClicked = (cS == LOW && lastS == HIGH);

  lastL = cL; lastR = cR; lastS = cS;
}

// --- SNAKE LOGIC ---
void initSnake() {
  snakeLen = 3; sDir = 0; snakeSpeed = 150; score = 0;
  for(uint8_t i=0; i<snakeLen; i++) snake[i] = { (int8_t)(10-i), 8 };
  food = { (int8_t)random(0,30), (int8_t)random(0,15) };
  gameRunning = true;
}

void runSnake() {
  // FIXED DIRECTIONS: 
  // If Left button was turning you Right, we reverse these:
  if (leftClicked)  { sDir = (sDir + 1) % 4; leftClicked = false; } // Now turns one way
  if (rightClicked) { sDir = (sDir + 3) % 4; rightClicked = false; } // Now turns the other
  
  if (gameRunning) startClicked = false; 

  if (gameRunning && millis() - prevTime > snakeSpeed) {
    prevTime = millis();
    for (int i = snakeLen-1; i > 0; i--) snake[i] = snake[i-1];
    snake[0].x += sdx[sDir]; 
    snake[0].y += sdy[sDir];

    if (snake[0].x < 0 || snake[0].x >= 32 || snake[0].y < 0 || snake[0].y >= 16) gameRunning = false;
    for (int i = 1; i < snakeLen; i++) 
      if (snake[0].x == snake[i].x && snake[0].y == snake[i].y) gameRunning = false;
    
    if (snake[0].x == food.x && snake[0].y == food.y) {
      if (snakeLen < 40) snakeLen++;
      score++;
      food = { (int8_t)random(0,31), (int8_t)random(0,15) };
      if (snakeSpeed > 60) snakeSpeed -= 2;
    }
  }

  u8g2.firstPage();
  do {
    if (!gameRunning) {
      u8g2.setFont(u8g2_font_6x10_tf); u8g2.drawStr(30, 32, "SNAKE OVER");
      if (startClicked) { startClicked = false; initSnake(); }
      if (digitalRead(BTN_L) == LOW && digitalRead(BTN_R) == LOW) currentState = MENU;
    } else {
      u8g2.drawBox(food.x*S_GRID, food.y*S_GRID, 3, 3);
      for (int i=0; i<snakeLen; i++) u8g2.drawFrame(snake[i].x*S_GRID, snake[i].y*S_GRID, 4, 4);
    }
  } while (u8g2.nextPage());
}

// --- PONG LOGIC ---
void initPong() {
  score = 0; paddleX = 54;
  ballX = 64; ballY = 20; ballDX = 1.5; ballDY = 1.5;
  gameRunning = true;
}

void runPong() {
  // Use raw digitalRead for smooth paddle movement
  if (digitalRead(BTN_L) == LOW && paddleX > 0) paddleX -= 3;
  if (digitalRead(BTN_R) == LOW && paddleX < 108) paddleX += 3;
  
  if (gameRunning) {
    ballX += ballDX; ballY += ballDY;
    if (ballX <= 0 || ballX >= 127) ballDX *= -1;
    if (ballY <= 0) ballDY *= -1;
    if (ballY >= 58 && ballX >= paddleX && ballX <= paddleX + 20) {
      ballDY *= -1;
      score++;
    }
    if (ballY > 64) gameRunning = false;
  }

  u8g2.firstPage();
  do {
    if (!gameRunning) {
      u8g2.setFont(u8g2_font_6x10_tf); u8g2.drawStr(30, 32, "PONG OVER");
      if (startClicked) { startClicked = false; initPong(); }
      if (digitalRead(BTN_L) == LOW && digitalRead(BTN_R) == LOW) currentState = MENU;
    } else {
      u8g2.drawBox(paddleX, 60, 20, 3);
      u8g2.drawDisc(ballX, ballY, 2);
    }
  } while (u8g2.nextPage());
}

// --- MENU ---
void handleMenu() {
  if (leftClicked) { menuSelection = (menuSelection + 2) % 3; leftClicked = false; }
  if (rightClicked) { menuSelection = (menuSelection + 1) % 3; rightClicked = false; }
  
  if (startClicked) {
    startClicked = false;
    if (menuSelection == 0) { initSnake(); currentState = SNAKE_GAME; }
    else if (menuSelection == 1) { initPong(); currentState = PONG_GAME; }
    // Timing logic placeholder
  }

  u8g2.firstPage();
  do {
    u8g2.setFont(u8g2_font_ncenB10_tr); u8g2.drawStr(35, 15, "GAMES");
    u8g2.setFont(u8g2_font_6x12_tr);
    u8g2.drawStr(15, 35, (menuSelection == 0) ? "> SNAKE" : "  SNAKE");
    u8g2.drawStr(15, 50, (menuSelection == 1) ? "> PONG" : "  PONG");
    u8g2.drawStr(15, 64, (menuSelection == 2) ? "> TIMING" : "  TIMING");
  } while (u8g2.nextPage());
}

void setup() {
  u8g2.begin();
  pinMode(BTN_L, INPUT_PULLUP); 
  pinMode(BTN_R, INPUT_PULLUP); 
  pinMode(BTN_START, INPUT_PULLUP);
  randomSeed(analogRead(A0));
}

void loop() {
  updateButtons();
  switch (currentState) {
    case MENU:        handleMenu(); break;
    case SNAKE_GAME:  runSnake();   break;
    case PONG_GAME:   runPong();    break;
  }
}