import pygame
import sys
import random

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

# --- 常量定义 ---
SCREEN_WIDTH = 600
SCREEN_HEIGHT = 600
GRID_SIZE = 20  # 网格大小
GRID_WIDTH = SCREEN_WIDTH // GRID_SIZE
GRID_HEIGHT = SCREEN_HEIGHT // GRID_SIZE
FPS = 10  # 游戏速度（帧率）

# 颜色定义 (R, G, B)
COLOR_BG = (20, 20, 30)       # 深蓝黑背景
COLOR_SNAKE = (0, 255, 150)   # 蛇身颜色 (亮绿)
COLOR_SNAKE_HEAD = (0, 200, 100) # 蛇头颜色 (深绿)
COLOR_FOOD = (255, 50, 50)    # 食物颜色 (红)
COLOR_TEXT = (255, 255, 255)  # 文字颜色
COLOR_GRID = (30, 30, 40)     # 网格线颜色

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

class Snake:
    def __init__(self):
        self.reset()

    def reset(self):
        # 初始位置在屏幕中心
        start_x = GRID_WIDTH // 2
        start_y = GRID_HEIGHT // 2
        self.body = [(start_x, start_y), (start_x - 1, start_y), (start_x - 2, start_y)]
        self.direction = RIGHT
        self.grow = False

    def move(self):
        head_x, head_y = self.body[0]
        dir_x, dir_y = self.direction
        new_head = (head_x + dir_x, head_y + dir_y)
        
        # 插入新头部
        self.body.insert(0, new_head)
        
        # 如果不吃食物，移除尾部
        if not self.grow:
            self.body.pop()
        else:
            self.grow = False

    def change_direction(self, new_dir):
        # 防止反向移动 (例如：向右走时不能直接向左)
        if (new_dir[0] * -1, new_dir[1] * -1) != self.direction:
            self.direction = new_dir

    def check_collision(self):
        head = self.body[0]
        
        # 撞墙检测
        if (head[0] < 0 or head[0] >= GRID_WIDTH or 
            head[1] < 0 or head[1] >= GRID_HEIGHT):
            return True
            
        # 撞自己检测
        if head in self.body[1:]:
            return True
            
        return False

    def draw(self, surface):
        for i, (x, y) in enumerate(self.body):
            rect = pygame.Rect(x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE)
            # 蛇头颜色稍深，蛇身颜色稍浅
            color = COLOR_SNAKE_HEAD if i == 0 else COLOR_SNAKE
            pygame.draw.rect(surface, color, rect)
            # 给蛇身加个边框，看起来更有质感
            pygame.draw.rect(surface, COLOR_BG, rect, 1)

class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize()

    def randomize(self, snake_body=None):
        while True:
            x = random.randint(0, GRID_WIDTH - 1)
            y = random.randint(0, GRID_HEIGHT - 1)
            self.position = (x, y)
            # 确保食物不会生成在蛇身上
            if snake_body is None or self.position not in snake_body:
                break

    def draw(self, surface):
        x, y = self.position
        rect = pygame.Rect(x * GRID_SIZE, y * GRID_SIZE, GRID_SIZE, GRID_SIZE)
        # 画个圆形的食物
        pygame.draw.ellipse(surface, COLOR_FOOD, rect)
        # 加点高光
        highlight = pygame.Rect(x * GRID_SIZE + 4, y * GRID_SIZE + 4, 4, 4)
        pygame.draw.ellipse(surface, (255, 200, 200), highlight)

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

def draw_text(surface, text, size, x, y, color=COLOR_TEXT, center=True):
    font = pygame.font.SysFont('Arial', size, bold=True)
    text_surface = font.render(text, True, color)
    text_rect = text_surface.get_rect()
    if center:
        text_rect.center = (x, y)
    else:
        text_rect.topleft = (x, y)
    surface.blit(text_surface, text_rect)

def main():
    screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
    pygame.display.set_caption("Python 贪吃蛇大作战")
    clock = pygame.time.Clock()

    snake = Snake()
    food = Food()
    food.randomize(snake.body)
    
    score = 0
    high_score = 0
    state = "START" # START, PLAYING, PAUSED, GAMEOVER

    running = True
    while running:
        # 1. 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            
            if event.type == pygame.KEYDOWN:
                if state == "START":
                    if event.key == pygame.K_SPACE:
                        state = "PLAYING"
                        snake.reset()
                        score = 0
                        food.randomize(snake.body)
                
                elif state == "PLAYING":
                    if event.key == pygame.K_UP or event.key == pygame.K_w:
                        snake.change_direction(UP)
                    elif event.key == pygame.K_DOWN or event.key == pygame.K_s:
                        snake.change_direction(DOWN)
                    elif event.key == pygame.K_LEFT or event.key == pygame.K_a:
                        snake.change_direction(LEFT)
                    elif event.key == pygame.K_RIGHT or event.key == pygame.K_d:
                        snake.change_direction(RIGHT)
                    elif event.key == pygame.K_ESCAPE:
                        state = "PAUSED"
                
                elif state == "PAUSED":
                    if event.key == pygame.K_ESCAPE:
                        state = "PLAYING"
                
                elif state == "GAMEOVER":
                    if event.key == pygame.K_SPACE:
                        state = "START"

        # 2. 逻辑更新
        if state == "PLAYING":
            snake.move()
            
            # 吃食物检测
            if snake.body[0] == food.position:
                snake.grow = True
                score += 1
                if score > high_score:
                    high_score = score
                food.randomize(snake.body)
            
            # 碰撞检测
            if snake.check_collision():
                state = "GAMEOVER"

        # 3. 绘制渲染
        screen.fill(COLOR_BG)
        draw_grid(screen) # 绘制背景网格

        # 绘制游戏元素
        if state != "START":
            food.draw(screen)
            snake.draw(screen)

        # 绘制 UI
        draw_text(screen, f"Score: {score}", 20, 10, 10, center=False)
        draw_text(screen, f"High Score: {high_score}", 20, SCREEN_WIDTH - 10, 10, center=False)

        # 状态界面
        if state == "START":
            draw_text(screen, "贪吃蛇大作战", 60, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 - 50)
            draw_text(screen, "按 [空格] 开始游戏", 30, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 20)
            draw_text(screen, "方向键/WASD 移动 | ESC 暂停", 20, SCREEN_WIDTH//2, SCREEN_HEIGHT - 40, (150, 150, 150))
        
        elif state == "PAUSED":
            # 半透明遮罩
            s = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
            s.set_alpha(150)
            s.fill((0, 0, 0))
            screen.blit(s, (0, 0))
            draw_text(screen, "游戏暂停", 50, SCREEN_WIDTH//2, SCREEN_HEIGHT//2)
            draw_text(screen, "按 [ESC] 继续", 25, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 50)

        elif state == "GAMEOVER":
            # 半透明遮罩
            s = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
            s.set_alpha(180)
            s.fill((0, 0, 0))
            screen.blit(s, (0, 0))
            draw_text(screen, "GAME OVER", 60, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 - 40, (255, 50, 50))
            draw_text(screen, f"最终得分: {score}", 40, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 20)
            draw_text(screen, "按 [空格] 返回主菜单", 25, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 70)

        pygame.display.flip()
        clock.tick(FPS)

    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()