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

// 屏幕尺寸定义 / Display dimensions
#define SCREEN_WIDTH        128 // OLED显示宽度(像素) / OLED display width, in pixels
#define SCREEN_HEIGHT       64  // OLED显示高度(像素) / OLED display height, in pixels

// OLED显示设置 / OLED display setup
#define OLED_RESET          4   // 复位引脚 / Reset pin # (or -1 if sharing Arduino reset pin)
Adafruit_SSD1306 display(SCREEN_WIDTH, SCREEN_HEIGHT, &Wire, OLED_RESET);

// 按钮引脚定义 (左, 上, 右, 下) / Button pins (LEFT, UP, RIGHT, DOWN)
const byte buttonPins[] = {4, 2, 5, 3};

// 游戏状态枚举 / Game state enumeration
typedef enum {
  START,      // 开始界面 / Start screen
  RUNNING,    // 游戏进行中 / Game running
  GAMEOVER    // 游戏结束 / Game over
} State;

// 移动方向枚举 / Direction enumeration
typedef enum {
    LEFT,      // 左 / Left
    UP,        // 上 / Up
    RIGHT,     // 右 / Right
    DOWN       // 下 / Down
} Direction;

// 游戏常量定义 / Game constants
#define SNAKE_PIECE_SIZE     3   // 蛇身每节大小 / Size of each snake segment
#define MAX_SANKE_LENGTH    165  // 蛇最大长度 / Maximum snake length
#define MAP_SIZE_X          20   // 地图宽度(格子数) / Map width (in grid units)
#define MAP_SIZE_Y          20   // 地图高度(格子数) / Map height (in grid units)
#define STARTING_SNAKE_SIZE 5    // 蛇起始长度 / Starting snake length
#define SNAKE_MOVE_DELAY    30   // 蛇移动延迟(帧数) / Snake movement delay (in frames)

// 全局变量 / Global variables
State gameState;                // 当前游戏状态 / Current game state
int8_t snake[MAX_SANKE_LENGTH][2]; // 蛇身体坐标数组 / Snake body coordinates array
uint8_t snake_length;           // 当前蛇长度 / Current snake length
Direction dir;                  // 当前移动方向 / Current movement direction
Direction newDir;               // 新设定的移动方向 / New movement direction
int8_t fruit[2];                // 食物坐标 / Fruit coordinates

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

  // 初始化OLED显示 / Initialize OLED display
  if(!display.begin(SSD1306_SWITCHCAPVCC, 0x3D)) {
    Serial.println(F("SSD1306分配失败 / SSD1306 allocation failed"));
    for(;;); // 卡住 / Hang
  }

  // 设置按钮引脚为输入上拉 / Set button pins as input pullup
  for (byte i = 0; i < 4; i++) {
    pinMode(buttonPins[i], INPUT_PULLUP);
  }

  // 初始化随机数种子 / Initialize random seed
  randomSeed(analogRead(A0));

  // 设置游戏 / Setup game
  setupGame();
}

// 初始化游戏设置 / Initialize game settings
void setupGame() {
  gameState = START;    // 初始状态为开始界面 / Initial state is start screen
  dir = RIGHT;          // 初始方向向右 / Initial direction is right
  newDir = RIGHT;       // 新方向也设为右 / New direction also right
  resetSnake();         // 重置蛇 / Reset snake
  generateFruit();      // 生成食物 / Generate fruit
  display.clearDisplay(); // 清屏 / Clear display
  drawMap();            // 绘制地图 / Draw map
  drawScore();          // 绘制分数 / Draw score
  drawPressToStart();   // 绘制开始提示 / Draw start prompt
  display.display();    // 更新显示 / Update display
}

// 重置蛇到初始位置和长度 / Reset snake to initial position and length
void resetSnake() {
  snake_length = STARTING_SNAKE_SIZE;
  for(int i = 0; i < snake_length; i++) {
    snake[i][0] = MAP_SIZE_X / 2 - i; // 水平居中 / Horizontal center
    snake[i][1] = MAP_SIZE_Y / 2;     // 垂直居中 / Vertical center
  }
}

int moveTime = 0; // 移动计时器 / Movement timer

void loop() {
  switch(gameState) {
    case START:
      // 开始界面：等待任意按钮按下 / Start screen: wait for any button press
      if(buttonPress()) gameState = RUNNING;
      break;
    
    case RUNNING:
      moveTime++; // 增加移动计时器 / Increment movement timer
      readDirection(); // 读取方向输入 / Read direction input
      
      // 按固定间隔移动蛇 / Move snake at fixed intervals
      if(moveTime >= SNAKE_MOVE_DELAY) {
        dir = newDir; // 应用新方向 / Apply new direction
        display.clearDisplay(); // 清屏 / Clear display
        
        // 移动蛇并检测碰撞 / Move snake and check for collisions
        if(moveSnake()) {
          gameState = GAMEOVER; // 如果碰撞，游戏结束 / If collision, game over
          drawGameover();       // 绘制游戏结束画面 / Draw game over screen
          delay(1000);          // 延迟1秒 / Delay 1 second
        }
        
        drawMap();    // 重绘地图 / Redraw map
        drawScore();  // 重绘分数 / Redraw score
        display.display(); // 更新显示 / Update display
        checkFruit(); // 检查是否吃到食物 / Check if fruit is eaten
        moveTime = 0; // 重置移动计时器 / Reset movement timer
      }
      break;
    
    case GAMEOVER:
      // 游戏结束：等待任意按钮按下重新开始 / Game over: wait for any button press to restart
      if(buttonPress()) {
        delay(500); // 防抖延迟 / Debounce delay
        setupGame(); // 重新设置游戏 / Reset game
        gameState = START; // 返回开始界面 / Back to start screen
      }
      break;
  }
  
  delay(10); // 短暂延迟 / Short delay
}

// 检测是否有按钮按下 / Check if any button is pressed
bool buttonPress() {
  for (byte i = 0; i < 4; i++) {
    if (digitalRead(buttonPins[i]) == LOW) {
      return true;
    }
  }
  return false;
}

// 读取方向输入（防止180度转弯） / Read direction input (prevent 180-degree turns)
void readDirection() {
  for (byte i = 0; i < 4; i++) {
    if (digitalRead(buttonPins[i]) == LOW && i != ((int)dir + 2) % 4) {
      newDir = (Direction)i;
      return;
    }
  }
}

// 移动蛇并检测碰撞 / Move snake and check for collisions
bool moveSnake() {
  int8_t x = snake[0][0]; // 蛇头X坐标 / Snake head X coordinate
  int8_t y = snake[0][1]; // 蛇头Y坐标 / Snake head Y coordinate

  // 根据方向计算新位置 / Calculate new position based on direction
  switch(dir) {
    case LEFT:  x -= 1; break;
    case UP:    y -= 1; break;
    case RIGHT: x += 1; break;
    case DOWN:  y += 1; break;
  }

  // 碰撞检测 / Collision detection
  if(collisionCheck(x, y))
    return true;

  // 移动蛇身（从尾部开始更新位置） / Move snake body (update positions from tail)
  for(int i = snake_length - 1; i > 0; i--) {
    snake[i][0] = snake[i - 1][0];
    snake[i][1] = snake[i - 1][1];
  }

  // 更新蛇头位置 / Update head position
  snake[0][0] = x;
  snake[0][1] = y;
  return false;
}

// 检查是否吃到食物 / Check if fruit is eaten
void checkFruit() {
  if(fruit[0] == snake[0][0] && fruit[1] == snake[0][1]) {
    // 增加蛇长度（不超过最大值） / Increase snake length (within limit)
    if(snake_length + 1 <= MAX_SANKE_LENGTH)
      snake_length++;
    generateFruit(); // 生成新食物 / Generate new fruit
  }
}

// 生成食物（确保不生成在蛇身上） / Generate fruit (ensure it's not on snake)
void generateFruit() {
  bool onSnake;
  do {
    onSnake = false;
    fruit[0] = random(0, MAP_SIZE_X); // 随机X坐标 / Random X coordinate
    fruit[1] = random(0, MAP_SIZE_Y); // 随机Y坐标 / Random Y coordinate
    
    // 检查是否与蛇身重叠 / Check if overlaps with snake
    for(int i = 0; i < snake_length; i++) {
      if(fruit[0] == snake[i][0] && fruit[1] == snake[i][1]) {
        onSnake = true;
        break;
      }
    }
  } while(onSnake);
}

// 碰撞检测 / Collision detection
bool collisionCheck(int8_t x, int8_t y) {
  // 检查是否撞到自己 / Check if hit itself
  for(int i = 1; i < snake_length; i++) {
    if(x == snake[i][0] && y == snake[i][1]) return true;
  }
  
  // 检查是否撞墙 / Check if hit wall
  if(x < 0 || y < 0 || x >= MAP_SIZE_X || y >= MAP_SIZE_Y) return true;
  return false;
}

// 绘制地图（蛇和食物） / Draw map (snake and fruit)
void drawMap() {
  // 计算地图偏移量（居中显示） / Calculate map offset (for centering)
  int offsetMapX = SCREEN_WIDTH - SNAKE_PIECE_SIZE * MAP_SIZE_X - 2;
  int offsetMapY = 2;

  // 绘制食物（反色显示） / Draw fruit (inverse color)
  display.drawRect(fruit[0] * SNAKE_PIECE_SIZE + offsetMapX, 
                  fruit[1] * SNAKE_PIECE_SIZE + offsetMapY, 
                  SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_INVERSE);
  
  // 绘制地图边框 / Draw map border
  display.drawRect(offsetMapX - 2, 0, 
                  SNAKE_PIECE_SIZE * MAP_SIZE_X + 4, 
                  SNAKE_PIECE_SIZE * MAP_SIZE_Y + 4, SSD1306_WHITE);
  
  // 绘制蛇身 / Draw snake body
  for(int i = 0; i < snake_length; i++) {
    display.fillRect(snake[i][0] * SNAKE_PIECE_SIZE + offsetMapX, 
                    snake[i][1] * SNAKE_PIECE_SIZE + offsetMapY, 
                    SNAKE_PIECE_SIZE, SNAKE_PIECE_SIZE, SSD1306_WHITE);
  }
}

// 绘制分数 / Draw score
void drawScore() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 2);
  display.print(F("Score:"));
  display.println(snake_length - STARTING_SNAKE_SIZE); // 分数=当前长度-起始长度
}

// 绘制开始提示 / Draw start prompt
void drawPressToStart() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 20);
  display.print(F("Press a\n button to\n start the\n game!"));
}

// 绘制游戏结束 / Draw game over
void drawGameover() {
  display.setTextSize(1);
  display.setTextColor(SSD1306_WHITE);
  display.setCursor(2, 50);
  display.println(F("GAMEOVER"));
}