import pygame
import random

# 初始化Pygame
pygame.init()

# 窗口与格子设置
WIDTH, HEIGHT = 640, 480
GRID = 20
speed = 12

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

# 创建窗口 & 时钟
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("经典贪吃蛇")
clock = pygame.time.Clock()

# 字体
font = pygame.font.SysFont(None, 36)

# 绘制分数


def draw_score(score):
    txt = font.render(f"Score: {score}", True, WHITE)
    screen.blit(txt, (10, 10))

# 游戏主逻辑


def game_run():
    # 蛇初始
    x, y = WIDTH//2, HEIGHT//2
    dx, dy = GRID, 0
    snake = [(x, y)]
    length = 1

    # 随机食物
    def new_food():
        fx = random.randrange(0, WIDTH-GRID, GRID)
        fy = random.randrange(0, HEIGHT-GRID, GRID)
        return fx, fy
    food = new_food()

    game_over = False
    while not game_over:
        # 事件监听
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                exit()
            # 方向控制（禁止直接反向）
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP and dy == 0:
                    dx, dy = 0, -GRID
                elif event.key == pygame.K_DOWN and dy == 0:
                    dx, dy = 0, GRID
                elif event.key == pygame.K_LEFT and dx == 0:
                    dx, dy = -GRID, 0
                elif event.key == pygame.K_RIGHT and dx == 0:
                    dx, dy = GRID, 0

        # 更新蛇头
        x += dx
        y += dy

        # 撞墙判定
        if x < 0 or x >= WIDTH or y < 0 or y >= HEIGHT:
            game_over = True
        # 撞自身
        if (x, y) in snake:
            game_over = True

        snake.append((x, y))
        if len(snake) > length:
            snake.pop(0)

        # 绘图
        screen.fill(BLACK)
        # 食物
        pygame.draw.rect(screen, RED, (food[0], food[1], GRID, GRID))
        # 蛇
        for seg in snake:
            pygame.draw.rect(screen, GREEN, (seg[0], seg[1], GRID, GRID))

        draw_score(length - 1)
        pygame.display.update()

        # 吃到食物
        if x == food[0] and y == food[1]:
            length += 1
            food = new_food()

        clock.tick(speed)

    # 结束提示 + 重启/退出
    while True:
        screen.fill(BLACK)
        end_text = font.render("Game Over! R重来 / Q退出", True, RED)
        screen.blit(end_text, (WIDTH//6, HEIGHT//2))
        pygame.display.update()
        for e in pygame.event.get():
            if e.type == pygame.QUIT:
                pygame.quit()
                exit()
            if e.type == pygame.KEYDOWN:
                if e.key == pygame.K_r:
                    game_run()
                if e.key == pygame.K_q:
                    pygame.quit()
                    exit()


# 启动游戏
if __name__ == "__main__":
    game_run()