import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏设置
WINDOW_WIDTH = 600
WINDOW_HEIGHT = 400
GRID_SIZE = 20
GRID_WIDTH = WINDOW_WIDTH // GRID_SIZE
GRID_HEIGHT = WINDOW_HEIGHT // GRID_SIZE

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 50, 50)
GREEN = (50, 200, 100)
DARK_GREEN = (30, 150, 70)
BLUE = (50, 130, 230)

class Snake:
    def __init__(self):
        self.length = 3
        self.positions = [(GRID_WIDTH // 2, GRID_HEIGHT // 2)]
        self.direction = (1, 0)  # 初始向右移动
        self.grow_flag = False
    
    def move(self):
        head = self.positions[0]
        new_head = (
            (head[0] + self.direction[0]) % GRID_WIDTH,
            (head[1] + self.direction[1]) % GRID_HEIGHT
        )
        
        if not self.grow_flag:
            self.positions.pop()
        else:
            self.grow_flag = False
        
        self.positions.insert(0, new_head)
    
    def grow(self):
        self.grow_flag = True
    
    def check_collision(self):
        return len(self.positions) != len(set(self.positions))
    
    def change_direction(self, direction):
        # 防止180度转弯
        if (direction[0] * -1, direction[1] * -1) != self.direction:
            self.direction = direction

class Food:
    def __init__(self):
        self.position = (0, 0)
        self.randomize_position()
    
    def randomize_position(self):
        self.position = (
            random.randint(0, GRID_WIDTH - 1),
            random.randint(0, GRID_HEIGHT - 1)
        )

def draw_grid(screen):
    for x in range(0, WINDOW_WIDTH, GRID_SIZE):
        pygame.draw.line(screen, (40, 45, 55), (x, 0), (x, WINDOW_HEIGHT))
    for y in range(0, WINDOW_HEIGHT, GRID_SIZE):
        pygame.draw.line(screen, (40, 45, 55), (0, y), (WINDOW_WIDTH, y))

def draw_snake(screen, snake):
    for i, pos in enumerate(snake.positions):
        rect = pygame.Rect(
            pos[0] * GRID_SIZE + 1,
            pos[1] * GRID_SIZE + 1,
            GRID_SIZE - 2,
            GRID_SIZE - 2
        )
        if i == 0:  # 蛇头
            pygame.draw.rect(screen, DARK_GREEN, rect)
            # 画眼睛
            eye_size = 4
            if snake.direction == (1, 0):  # 右
                eye_positions = [(pos[0] * GRID_SIZE + 14, pos[1] * GRID_SIZE + 5),
                                (pos[0] * GRID_SIZE + 15, pos[1] * GRID_SIZE + 13)]
            elif snake.direction == (-1, 0):  # 左
                eye_positions = [(pos[0] * GRID_SIZE + 5, pos[1] * GRID_SIZE + 5),
                                (pos[0] * GRID_SIZE + 6, pos[1] * GRID_SIZE + 8)]
            elif snake.direction == (0, -1):  # 上
                eye_positions = [(pos[0] * GRID_SIZE + 7, pos[1] * GRID_SIZE + 4),
                                (pos[0] * GRID_SIZE + 12, pos[1] * GRID_SIZE + 5)]
            else:  # 下
                eye_positions = [(pos[0] * GRID_SIZE + 7, pos[1] * GRID_SIZE + 14),
                                (pos[0] * GRID_SIZE + 11, pos[1] * GRID_SIZE + 16)]
            
            for eye_pos in eye_positions:
                pygame.draw.circle(screen, WHITE, eye_pos, eye_size // 2)
                pygame.draw.circle(screen, BLACK, eye_pos, eye_size // 4)
        else:  # 蛇身
            color = GREEN if i % 2 == 0 else DARK_GREEN
            pygame.draw.rect(screen, color, rect, border_radius=3)

def draw_food(screen, food):
    rect = pygame.Rect(
        food.position[0] * GRID_SIZE + 2,
        food.position[1] * GRID_SIZE + 2,
        GRID_SIZE - 4,
        GRID_SIZE - 4
    )
    pygame.draw.rect(screen, RED, rect, border_radius=5)
    # 添加高光效果
    highlight = pygame.Rect(
        food.position[0] * GRID_SIZE + 5,
        food.position[1] * GRID_SIZE + 5,
        GRID_SIZE // 3,
        GRID_SIZE // 3
    )
    pygame.draw.rect(screen, (255, 120, 80), highlight, border_radius=2)

def show_game_over(screen, score):
    font_large = pygame.font.Font(None, 72)
    font_small = pygame.font.Font(None, 36)
    
    game_over_text = font_large.render("Game Over!", True, RED)
    score_text = font_small.render(f"Score: {score}", True, WHITE)
    restart_text = font_small.render("Press SPACE to restart", True, BLUE)
    
    screen.blit(game_over_text, 
                (WINDOW_WIDTH // 2 - game_over_text.get_width() // 2, 
                 WINDOW_HEIGHT // 2 - 60))
    screen.blit(score_text, 
                (WINDOW_WIDTH // 2 - score_text.get_width() // 2, 
                 WINDOW_HEIGHT // 2))
    screen.blit(restart_text, 
                (WINDOW_WIDTH // 2 - restart_text.get_width() // 2, 
                 WINDOW_HEIGHT // 2 + 50))

def main():
    screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
    pygame.display.set_caption("贪吃蛇")
    clock = pygame.time.Clock()
    
    # 游戏变量
    snake = Snake()
    food = Food()
    score = 0
    game_over = False
    speed = 10
    
    # 字体
    font = pygame.font.Font(None, 24)
    
    while True:
        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 = Snake()
                        food = Food()
                        score = 0
                        game_over = False
                        speed = 10
                else:
                    if event.key == pygame.K_UP:
                        snake.change_direction((0, -1))
                    elif event.key == pygame.K_DOWN:
                        snake.change_direction((0, 1))
                    elif event.key == pygame.K_LEFT:
                        snake.change_direction((-1, 0))
                    elif event.key == pygame.K_RIGHT:
                        snake.change_direction((1, 0))
        
        if not game_over:
            # 更新蛇的位置
            snake.move()
            
            # 检查是否吃到食物
            if snake.positions[0] == food.position:
                snake.grow()
                score += 10
                food.randomize_position()
                # 每得50分加速一次
                if score % 50 == 0 and speed < 25:
                    speed += 2
            
            # 检查碰撞
            if snake.check_collision():
                game_over = True
        
        # 绘制画面
        screen.fill(BLACK)
        draw_grid(screen)
        draw_snake(screen, snake)
        draw_food(screen, food)
        
        # 显示分数
        score_text = font.render(f"Score: {score}", True, WHITE)
        screen.blit(score_text, (10, 10))
        
        if game_over:
            show_game_over(screen, score)
        
        pygame.display.flip()
        clock.tick(speed)

if __name__ == "__main__":
    main()