import pygame
import random
import sys

# 初始化
pygame.init()

# 窗口设置
WIDTH, HEIGHT = 1000, 800
BLOCK_SIZE = 20
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("贪吃蛇 - 元宝版")

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

# 时钟
clock = pygame.time.Clock()
speed = 8

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


def draw_text(text, color, x, y):
    img = font.render(text, True, color)
    screen.blit(img, (x, y))


def game_loop():
    # 蛇的初始位置
    snake = [(100, 100), (80, 100), (60, 100)]
    direction = (BLOCK_SIZE, 0)

    # 食物
    food = (
        random.randrange(0, WIDTH, BLOCK_SIZE),
        random.randrange(0, HEIGHT, BLOCK_SIZE),
    )

    score = 0

    running = True
    while running:
        clock.tick(speed)

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_UP and direction != (0, BLOCK_SIZE):
                    direction = (0, -BLOCK_SIZE)
                elif event.key == pygame.K_DOWN and direction != (0, -BLOCK_SIZE):
                    direction = (0, BLOCK_SIZE)
                elif event.key == pygame.K_LEFT and direction != (BLOCK_SIZE, 0):
                    direction = (-BLOCK_SIZE, 0)
                elif event.key == pygame.K_RIGHT and direction != (-BLOCK_SIZE, 0):
                    direction = (BLOCK_SIZE, 0)

        # 移动蛇
        head_x = snake[0][0] + direction[0]
        head_y = snake[0][1] + direction[1]
        new_head = (head_x, head_y)

        # 碰撞检测
        if (
            head_x < 0 or head_x >= WIDTH or
            head_y < 0 or head_y >= HEIGHT or
            new_head in snake
        ):
            running = False

        snake.insert(0, new_head)

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

        # 绘制
        screen.fill(BLACK)

        for block in snake:
            pygame.draw.rect(
                screen, GREEN,
                (block[0], block[1], BLOCK_SIZE, BLOCK_SIZE)
            )

        pygame.draw.rect(
            screen, RED,
            (food[0], food[1], BLOCK_SIZE, BLOCK_SIZE)
        )

        draw_text(f"Score: {score}", WHITE, 10, 10)

        pygame.display.flip()

    # 游戏结束画面
    screen.fill(BLACK)
    draw_text("Game Over", RED, WIDTH // 2 - 80, HEIGHT // 2 - 20)
    draw_text(f"Score: {score}", WHITE, WIDTH // 2 - 70, HEIGHT // 2 + 20)
    pygame.display.flip()
    pygame.time.delay(2000)


if __name__ == "__main__":
    game_loop()