import pygame
import sys
import random

# 初始化pygame
pygame.init()

# 游戏常量定义
WINDOW_WIDTH = 800
WINDOW_HEIGHT = 600
BLOCK_SIZE = 20  # 蛇和食物的方块大小
SPEED = 10       # 游戏速度（帧率）

# 颜色定义 (RGB)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)    # 食物颜色
GREEN = (0, 255, 0)  # 蛇身颜色
BLUE = (0, 0, 255)   # 蛇头颜色
GRAY = (128, 128, 128)  # 边框颜色

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

class Snake:
    """贪吃蛇类"""
    def __init__(self):
        # 初始位置：屏幕中心
        self.x = WINDOW_WIDTH // 2
        self.y = WINDOW_HEIGHT // 2
        # 初始方向：向右
        self.direction = RIGHT
        # 蛇身（列表存储每个方块的坐标，第一个元素是蛇头）
        self.body = [(self.x, self.y), 
                     (self.x - BLOCK_SIZE, self.y), 
                     (self.x - 2*BLOCK_SIZE, self.y)]
        # 是否吃到食物
        self.grow = False

    def move(self):
        """移动蛇（更新坐标）"""
        # 计算新蛇头坐标
        head_x = self.body[0][0] + self.direction[0]
        head_y = self.body[0][1] + self.direction[1]
        new_head = (head_x, head_y)
        
        # 添加新蛇头到身体最前面
        self.body.insert(0, new_head)
        
        # 如果没吃到食物，删除最后一个方块（保持长度）
        if not self.grow:
            self.body.pop()
        else:
            self.grow = False  # 重置增长标记

    def change_direction(self, new_dir):
        """改变移动方向（防止180度掉头）"""
        # 不能直接反向移动（比如向右时不能直接向左）
        if (new_dir[0] != -self.direction[0]) or (new_dir[1] != -self.direction[1]):
            self.direction = new_dir

    def check_collision(self):
        """检测碰撞（撞墙/撞自己）"""
        head_x, head_y = self.body[0]
        
        # 撞墙检测
        if (head_x < 0 or head_x >= WINDOW_WIDTH or
            head_y < 0 or head_y >= WINDOW_HEIGHT):
            return True
        
        # 撞自己检测
        if self.body[0] in self.body[1:]:
            return True
        
        return False

    def eat_food(self, food_pos):
        """检测是否吃到食物"""
        if self.body[0] == food_pos:
            self.grow = True
            return True
        return False

class Food:
    """食物类"""
    def __init__(self):
        # 随机生成食物位置（必须是BLOCK_SIZE的整数倍）
        self.x = random.randint(0, (WINDOW_WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
        self.y = random.randint(0, (WINDOW_HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE

    def respawn(self):
        """重新生成食物位置"""
        self.x = random.randint(0, (WINDOW_WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
        self.y = random.randint(0, (WINDOW_HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE

def main():
    """游戏主函数"""
    # 创建游戏窗口
    screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    pygame.display.set_caption("贪吃蛇游戏")
    
    # 创建时钟（控制帧率）
    clock = pygame.time.Clock()
    
    # 初始化字体（显示分数）
    font = pygame.font.Font(None, 40)
    
    # 创建蛇和食物对象
    snake = Snake()
    food = Food()
    
    # 游戏分数
    score = 0
    # 游戏结束标记
    game_over = False

    # 游戏主循环
    while True:
        # 事件处理
        for event in pygame.event.get():
            # 关闭窗口
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            
            # 键盘控制方向
            if event.type == pygame.KEYDOWN and not game_over:
                if event.key == pygame.K_UP:
                    snake.change_direction(UP)
                elif event.key == pygame.K_DOWN:
                    snake.change_direction(DOWN)
                elif event.key == pygame.K_LEFT:
                    snake.change_direction(LEFT)
                elif event.key == pygame.K_RIGHT:
                    snake.change_direction(RIGHT)
            
            # 游戏结束后按空格键重新开始
            if event.type == pygame.KEYDOWN and game_over:
                if event.key == pygame.K_SPACE:
                    # 重置游戏状态
                    snake = Snake()
                    food = Food()
                    score = 0
                    game_over = False

        if not game_over:
            # 移动蛇
            snake.move()
            
            # 检测碰撞（游戏结束）
            if snake.check_collision():
                game_over = True
            
            # 检测是否吃到食物
            if snake.eat_food((food.x, food.y)):
                score += 10  # 每吃一个食物加10分
                food.respawn()  # 重新生成食物

        # 绘制界面
        screen.fill(BLACK)
        
        # 绘制蛇（蛇头用蓝色，身体用绿色）
        for i, (x, y) in enumerate(snake.body):
            color = BLUE if i == 0 else GREEN
            pygame.draw.rect(screen, color, (x, y, BLOCK_SIZE, BLOCK_SIZE))
            # 给蛇块加边框（更清晰）
            pygame.draw.rect(screen, GRAY, (x, y, BLOCK_SIZE, BLOCK_SIZE), 1)
        
        # 绘制食物
        pygame.draw.rect(screen, RED, (food.x, food.y, BLOCK_SIZE, BLOCK_SIZE))
        pygame.draw.rect(screen, GRAY, (food.x, food.y, BLOCK_SIZE, BLOCK_SIZE), 1)
        
        # 绘制分数
        score_text = font.render(f"Score: {score}", True, WHITE)
        screen.blit(score_text, (10, 10))
        
        # 游戏结束提示
        if game_over:
            game_over_text = font.render("Game Over! Press SPACE to restart", True, RED)
            text_rect = game_over_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2))
            screen.blit(game_over_text, text_rect)
        
        # 更新屏幕显示
        pygame.display.flip()
        
        # 控制游戏速度
        clock.tick(SPEED)

if __name__ == "__main__":
    main()