#include "stm32c0xx.h"
#include <Wire.h>
#include <Adafruit_GFX.h>
#include <Adafruit_SSD1306.h>
#include <stdlib.h>
#include <time.h>

#define WIDTH 128
#define HEIGHT 64
#define OLED_RESET -1
#define GRID_SIZE 4
#define GRID_WIDTH (WIDTH / GRID_SIZE)
#define GRID_HEIGHT (HEIGHT / GRID_SIZE)

typedef struct {
    int x;
    int y;
} Point;
Adafruit_SSD1306 display(WIDTH, HEIGHT, &Wire, OLED_RESET);
Point snake[100];
int snake_length = 3;
Point food;
int direction = 0; // 0: right, 1: down, 2: left, 3: up
int game_over = 0;
int score = 0;

void generate_food() {
    int valid;
    do {
        valid = 1;
        food.x = rand() % GRID_WIDTH;
        food.y = rand() % GRID_HEIGHT;
        
        // Check if food is on snake
        for (int i = 0; i < snake_length; i++) {
            if (snake[i].x == food.x && snake[i].y == food.y) {
                valid = 0;
                break;
            }
        }
    } while (!valid);
}

void init_game() {
    snake_length = 3;
    direction = 0;
    game_over = 0;
    score = 0;
    
    // Initialize snake in the middle
    snake[0].x = GRID_WIDTH / 2;
    snake[0].y = GRID_HEIGHT / 2;
    snake[1].x = snake[0].x - 1;
    snake[1].y = snake[0].y;
    snake[2].x = snake[0].x - 2;
    snake[2].y = snake[0].y;
    
    generate_food();
}

void update_game() {
    if (game_over) return;
    
    // Move snake
    Point new_head = snake[0];
    
    switch (direction) {
        case 0: new_head.x++; break; // right
        case 1: new_head.y++; break; // down
        case 2: new_head.x--; break; // left
        case 3: new_head.y--; break; // up
    }
    
    // Check collision with walls
    if (new_head.x < 0 || new_head.x >= GRID_WIDTH || 
        new_head.y < 0 || new_head.y >= GRID_HEIGHT) {
        game_over = 1;
        return;
    }
    
    // Check collision with self
    for (int i = 0; i < snake_length; i++) {
        if (snake[i].x == new_head.x && snake[i].y == new_head.y) {
            game_over = 1;
            return;
        }
    }
    
    // Move snake body
    for (int i = snake_length - 1; i > 0; i--) {
        snake[i] = snake[i - 1];
    }
    snake[0] = new_head;
    
    // Check if snake ate food
    if (snake[0].x == food.x && snake[0].y == food.y) {
        snake_length++;
        score += 10;
        generate_food();
        
        // Add new segment (it will be updated in next move)
        snake[snake_length - 1] = snake[snake_length - 2];
    }
}

void draw_game() {
    display.clearDisplay();
    
    if (game_over) {
        display.setTextSize(2);
        display.setCursor(10, 20);
        display.print("Game Over!");
        
        display.setTextSize(1);
        display.setCursor(20, 40);
        display.print("Score: ");
        display.print(score);
    } else {
        // 绘制食物（空心矩形）
        display.drawRect(
            food.x * GRID_SIZE,
            food.y * GRID_SIZE,
            GRID_SIZE,
            GRID_SIZE,
            SSD1306_WHITE
        );
        
        // 绘制蛇身（实心矩形）
        for (int i = 0; i < snake_length; i++) {
            display.fillRect(
                snake[i].x * GRID_SIZE,
                snake[i].y * GRID_SIZE,
                GRID_SIZE,
                GRID_SIZE,
                SSD1306_WHITE
            );
        }
        
        // 显示分数
        display.setTextSize(1);
        display.setCursor(0, 0);
        display.print("Score: ");
        display.print(score);
    }
    
    display.display();
}

void check_input() {
    static int last_btn_state[4] = {1, 1, 1, 1};
    int btn_state[4];
    
    // Read buttons (active low)
    btn_state[0] = !(GPIOA->IDR & (1 << 0)); // Up
    btn_state[1] = !(GPIOA->IDR & (1 << 1)); // Right
    btn_state[2] = !(GPIOA->IDR & (1 << 2)); // Down
    btn_state[3] = !(GPIOA->IDR & (1 << 3)); // Left
    
    // Check for button presses
    for (int i = 0; i < 4; i++) {
        if (btn_state[i] && !last_btn_state[i]) {
            // Prevent 180-degree turns
            if ((direction == 0 && i != 2) || 
                (direction == 1 && i != 3) ||
                (direction == 2 && i != 0) ||
                (direction == 3 && i != 1)) {
                direction = i;
            }
            if (game_over) {
                init_game();
            }
        }
        last_btn_state[i] = btn_state[i];
    }
}

void delay_ms(uint32_t ms) {
    for (uint32_t i = 0; i < ms * 1000; i++) {
        __asm__("nop");
    }
}

int main(void) {
    // 初始化按钮GPIO
    RCC->IOPENR |= RCC_IOPENR_GPIOAEN;
    GPIOA->MODER &= ~(GPIO_MODER_MODE0 | GPIO_MODER_MODE1 | GPIO_MODER_MODE2 | GPIO_MODER_MODE3);
    
    // 初始化I2C和OLED
    Wire.begin();
    display.clearDisplay();
    display.display();
    
    // 初始化随机数种子
    srand(HAL_GetTick());
    
    init_game();
    
    uint32_t last_update = 0;
    while (1) {
        check_input();
        
        uint32_t now = HAL_GetTick();
        if (now - last_update > 200) { // Update every 200ms
            update_game();
            draw_game();
            last_update = now;
        }
    }
}