import pygame
import random
import sys

# --- 初始化 Pygame ---
pygame.init()

# --- 常量定义 ---
SCREEN_WIDTH = 600
SCREEN_HEIGHT = 400
GRID_SIZE = 20
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // GRID_SIZE
INITIAL_FPS = 10

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
DARK_GREEN = (0, 200, 0)
RED = (255, 0, 0)
GRAY = (40, 40, 40)

# 方向向量
UP = (0, -1)
DOWN = (0, 1)
LEFT = (-1, 0)
RIGHT = (1, 0)

screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Pygame 贪吃蛇 (修复版)")
clock = pygame.time.Clock()
font = pygame.font.SysFont("Arial", 24)
large_font = pygame.font.SysFont("Arial", 48)


def draw_grid():
    for x in range(0, SCREEN_WIDTH, GRID_SIZE):
        pygame.draw.line(screen, GRAY, (x, 0), (x, SCREEN_HEIGHT))
    for y in range(0, SCREEN_HEIGHT, GRID_SIZE):
        pygame.draw.line(screen, GRAY, (0, y), (SCREEN_WIDTH, y))


def get_random_food(snake_body):
    """生成一个不在蛇身上的随机食物位置"""
    while True:
        fx = random.randrange(GRID_WIDTH) * GRID_SIZE
        fy = random.randrange(GRID_HEIGHT) * GRID_SIZE
        if (fx, fy) not in snake_body:
            return (fx, fy)


def show_game_over(score):
    screen.fill(BLACK)
    text_go = large_font.render("GAME OVER", True, RED)
    text_score = font.render(f"Final Score: {score}", True, WHITE)
    text_restart = font.render(
        "Press SPACE to Restart or Q to Quit", True, WHITE)

    rect_go = text_go.get_rect(center=(SCREEN_WIDTH/2, SCREEN_HEIGHT/2 - 50))
    rect_score = text_score.get_rect(
        center=(SCREEN_WIDTH/2, SCREEN_HEIGHT/2 + 10))
    rect_restart = text_restart.get_rect(
        center=(SCREEN_WIDTH/2, SCREEN_HEIGHT/2 + 60))

    screen.blit(text_go, rect_go)
    screen.blit(text_score, rect_score)
    screen.blit(text_restart, rect_restart)
    pygame.display.flip()

    waiting = True
    while waiting:
        clock.tick(INITIAL_FPS)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            if event.type == pygame.KEYUP:
                if event.key == pygame.K_q:
                    return False
                if event.key == pygame.K_SPACE:
                    return True
    return False


def main():
    running = True
    game_over = False

    # 游戏状态变量
    snake = []
    direction = RIGHT
    food = None
    score = 0
    current_fps = INITIAL_FPS

    # 初始化第一局游戏
    def reset_game():
        nonlocal snake, direction, food, score, current_fps, game_over
        snake = [(SCREEN_WIDTH//2, SCREEN_HEIGHT//2)]
        direction = RIGHT
        score = 0
        current_fps = INITIAL_FPS
        game_over = False
        food = get_random_food(snake)

    reset_game()

    while running:
        if game_over:
            if not show_game_over(score):
                running = False
                continue
            reset_game()

        # --- 事件处理 ---
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            if event.type == pygame.KEYDOWN and not game_over:
                if event.key == pygame.K_UP and direction != DOWN:
                    direction = UP
                elif event.key == pygame.K_DOWN and direction != UP:
                    direction = DOWN
                elif event.key == pygame.K_LEFT and direction != RIGHT:
                    direction = LEFT
                elif event.key == pygame.K_RIGHT and direction != LEFT:
                    direction = RIGHT

        if not game_over:
            # --- 游戏逻辑更新 ---
            head_x, head_y = snake[0]
            dir_x, dir_y = direction
            new_head = (head_x + dir_x * GRID_SIZE, head_y + dir_y * GRID_SIZE)

            # 1. 检查墙壁碰撞
            if (new_head[0] < 0 or new_head[0] >= SCREEN_WIDTH or
                    new_head[1] < 0 or new_head[1] >= SCREEN_HEIGHT):
                game_over = True

            # 2. 检查自身碰撞
            elif new_head in snake:
                game_over = True

            else:
                # 移动蛇
                snake.insert(0, new_head)

                # 3. 检查是否吃到食物
                # 此时 food 一定不为 None，因为 reset_game 保证了这一点
                if new_head == food:
                    score += 10
                    current_fps = min(INITIAL_FPS + (score // 50), 25)
                    food = get_random_food(snake)  # 生成新食物
                else:
                    snake.pop()  # 没吃到，移除尾巴

        # --- 绘制渲染 (无论是否 game_over 都要绘制最后一帧，除非刚重置) ---
        screen.fill(BLACK)
        draw_grid()

        # 画食物 (确保 food 存在)
        if food:
            food_rect = pygame.Rect(food[0], food[1], GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(screen, RED, food_rect)

        # 画蛇
        for i, segment in enumerate(snake):
            color = DARK_GREEN if i == 0 else GREEN
            rect = pygame.Rect(segment[0], segment[1], GRID_SIZE, GRID_SIZE)
            pygame.draw.rect(screen, color, rect)
            pygame.draw.rect(screen, BLACK, rect, 1)  # 边框

        # 画分数
        score_text = font.render(f"Score: {score}", True, WHITE)
        screen.blit(score_text, (10, 10))

        pygame.display.flip()

        # 控制速度
        clock.tick(current_fps if not game_over else INITIAL_FPS)

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()