import pygame
import sys
import random

pygame.init()

# 窗口设置
WIDTH, HEIGHT = 600, 400
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇")

# 颜色
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)

# 格子大小
BLOCK_SIZE = 20
SPEED = 10

# 字体（兼容所有系统）
font = pygame.font.Font(None, 35)

def reset_game():
    # 蛇初始位置
    snake = [
        [WIDTH//2, HEIGHT//2],
        [WIDTH//2 - BLOCK_SIZE, HEIGHT//2],
        [WIDTH//2 - BLOCK_SIZE*2, HEIGHT//2]
    ]
    dir_x = BLOCK_SIZE
    dir_y = 0
    food = [
        random.randrange(0, WIDTH - BLOCK_SIZE, BLOCK_SIZE),
        random.randrange(0, HEIGHT - BLOCK_SIZE, BLOCK_SIZE)
    ]
    score = 0
    return snake, dir_x, dir_y, food, score

snake, dir_x, dir_y, food, score = reset_game()
clock = pygame.time.Clock()
game_over = False

while True:
    screen.fill(BLACK)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        # 按键监听
        if event.type == pygame.KEYDOWN:
            if game_over:
                if event.key == pygame.K_SPACE:
                    snake, dir_x, dir_y, food, score = reset_game()
                    game_over = False
            else:
                # 防止直接反向撞死自己
                if event.key == pygame.K_UP and dir_y != BLOCK_SIZE:
                    dir_x = 0
                    dir_y = -BLOCK_SIZE
                elif event.key == pygame.K_DOWN and dir_y != -BLOCK_SIZE:
                    dir_x = 0
                    dir_y = BLOCK_SIZE
                elif event.key == pygame.K_LEFT and dir_x != BLOCK_SIZE:
                    dir_x = -BLOCK_SIZE
                    dir_y = 0
                elif event.key == pygame.K_RIGHT and dir_x != -BLOCK_SIZE:
                    dir_x = BLOCK_SIZE
                    dir_y = 0

    if not game_over:
        # 新蛇头
        head_x, head_y = snake[0]
        new_head = [head_x + dir_x, head_y + dir_y]
        snake.insert(0, new_head)

        # 吃到食物
        if new_head[0] == food[0] and new_head[1] == food[1]:
            score += 10
            food = [
                random.randrange(0, WIDTH - BLOCK_SIZE, BLOCK_SIZE),
                random.randrange(0, HEIGHT - BLOCK_SIZE, BLOCK_SIZE)
            ]
        else:
            snake.pop()

        # 撞墙检测
        if (new_head[0] < 0 or new_head[0] >= WIDTH or
            new_head[1] < 0 or new_head[1] >= HEIGHT):
            game_over = True

        # 撞到自己
        if new_head in snake[1:]:
            game_over = True

        # 绘制蛇
        for seg in snake:
            pygame.draw.rect(screen, GREEN, (seg[0], seg[1], BLOCK_SIZE-1, BLOCK_SIZE-1))
        # 绘制食物
        pygame.draw.rect(screen, RED, (food[0], food[1], BLOCK_SIZE-1, BLOCK_SIZE-1))

        # 分数
        score_text = font.render(f"分数：{score}", True, WHITE)
        screen.blit(score_text, (10, 10))
    else:
        # 游戏结束文字
        over_text = font.render("游戏结束！空格重新开始", True, (255,80,80))
        screen.blit(over_text, (WIDTH//2 - 220, HEIGHT//2))

    pygame.display.update()
    clock.tick(SPEED)