import pygame
import random
import sys

# 初始化Pygame
pygame.init()

# 游戏常量设置
WIDTH, HEIGHT = 800, 600
BLOCK_SIZE = 20
FPS = 10  # 基础帧率
FONT = pygame.font.Font(None, 40)

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
YELLOW = (255, 255, 0)  # 加速技能
PURPLE = (128, 0, 128)  # 穿墙技能
CYAN = (0, 255, 255)    # 变长技能

# 创建游戏窗口
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("技能贪吃蛇")
clock = pygame.time.Clock()

# 技能类：管理技能的冷却、持续时间和效果
class Skill:
    def __init__(self, cooldown, duration):
        self.cooldown = cooldown  # 技能冷却时间（帧）
        self.duration = duration  # 技能持续时间（帧）
        self.current_cooldown = 0  # 当前冷却剩余
        self.current_duration = 0  # 当前持续剩余
        self.active = False        # 技能是否激活

    def update(self):
        """每帧更新技能状态"""
        # 处理冷却
        if self.current_cooldown > 0:
            self.current_cooldown -= 1
        # 处理持续时间
        if self.active:
            self.current_duration -= 1
            if self.current_duration <= 0:
                self.active = False

    def use(self):
        """使用技能（需检查冷却）"""
        if self.current_cooldown == 0:
            self.active = True
            self.current_duration = self.duration
            self.current_cooldown = self.cooldown
            return True
        return False

# 贪吃蛇游戏主类
class SkillSnakeGame:
    def __init__(self):
        # 蛇的初始状态
        self.snake = [(WIDTH//2, HEIGHT//2), 
                      (WIDTH//2 - BLOCK_SIZE, HEIGHT//2),
                      (WIDTH//2 - 2*BLOCK_SIZE, HEIGHT//2)]
        self.direction = (BLOCK_SIZE, 0)  # 初始向右
        self.score = 0
        
        # 食物生成
        self.food = self.generate_food()
        self.skill_food = None  # 技能食物
        self.skill_food_timer = 0  # 技能食物生成计时器
        
        # 初始化技能
        self.speed_skill = Skill(cooldown=60, duration=30)  # 冷却60帧，持续30帧
        self.wall_skill = Skill(cooldown=120, duration=20)  # 冷却120帧，持续20帧
        self.length_skill = Skill(cooldown=90, duration=1)  # 冷却90帧，持续1帧（立即生效）
        
        # 游戏状态
        self.game_over = False
        self.paused = False

    def generate_food(self):
        """生成普通食物（避免生成在蛇身上）"""
        while True:
            x = random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            y = random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            food_pos = (x, y)
            if food_pos not in self.snake:
                return food_pos

    def generate_skill_food(self):
        """生成技能食物（随机类型）"""
        skill_types = ["speed", "wall", "length"]
        skill_type = random.choice(skill_types)
        while True:
            x = random.randint(0, (WIDTH - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            y = random.randint(0, (HEIGHT - BLOCK_SIZE) // BLOCK_SIZE) * BLOCK_SIZE
            skill_pos = (x, y)
            if skill_pos not in self.snake and skill_pos != self.food:
                return (skill_pos, skill_type)

    def handle_events(self):
        """处理键盘事件"""
        for event in pygame.event.get():
            # 关闭窗口
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            # 键盘按键
            if event.type == pygame.KEYDOWN:
                # 游戏结束后按R重启
                if self.game_over:
                    if event.key == pygame.K_r:
                        self.__init__()
                    continue
                # 暂停/继续
                if event.key == pygame.K_p:
                    self.paused = not self.paused
                # 方向控制（避免反向）
                if not self.paused:
                    if event.key == pygame.K_UP and self.direction != (0, BLOCK_SIZE):
                        self.direction = (0, -BLOCK_SIZE)
                    elif event.key == pygame.K_DOWN and self.direction != (0, -BLOCK_SIZE):
                        self.direction = (0, BLOCK_SIZE)
                    elif event.key == pygame.K_LEFT and self.direction != (BLOCK_SIZE, 0):
                        self.direction = (-BLOCK_SIZE, 0)
                    elif event.key == pygame.K_RIGHT and self.direction != (-BLOCK_SIZE, 0):
                        self.direction = (BLOCK_SIZE, 0)
                    # 技能使用
                    elif event.key == pygame.K_1:  # 1键：加速
                        self.speed_skill.use()
                    elif event.key == pygame.K_2:  # 2键：穿墙
                        self.wall_skill.use()
                    elif event.key == pygame.K_3:  # 3键：变长
                        if self.length_skill.use():
                            # 变长技能：直接增加3节身体
                            for _ in range(3):
                                self.snake.append(self.snake[-1])
                            self.score += 3

    def update_game(self):
        """更新游戏状态"""
        if self.game_over or self.paused:
            return
        
        # 更新技能状态
        self.speed_skill.update()
        self.wall_skill.update()
        self.length_skill.update()
        
        # 生成技能食物（每100帧随机生成）
        self.skill_food_timer += 1
        if self.skill_food_timer >= 100 and not self.skill_food:
            self.skill_food = self.generate_skill_food()
            self.skill_food_timer = 0
        
        # 移动蛇头
        head_x, head_y = self.snake[0]
        new_head = (head_x + self.direction[0], head_y + self.direction[1])
        
        # 穿墙技能处理
        if self.wall_skill.active:
            # 穿墙：超出边界则从对面出现
            new_head = (new_head[0] % WIDTH, new_head[1] % HEIGHT)
        else:
            # 无穿墙：撞墙游戏结束
            if (new_head[0] < 0 or new_head[0] >= WIDTH or
                new_head[1] < 0 or new_head[1] >= HEIGHT):
                self.game_over = True
        
        # 撞自己身体（游戏结束）
        if new_head in self.snake[:-1]:
            self.game_over = True
        
        # 添加新蛇头
        self.snake.insert(0, new_head)
        
        # 吃到普通食物
        if new_head == self.food:
            self.score += 1
            self.food = self.generate_food()
        # 吃到技能食物
        elif self.skill_food and new_head == self.skill_food[0]:
            skill_type = self.skill_food[1]
            if skill_type == "speed":
                self.speed_skill.use()  # 激活加速
            elif skill_type == "wall":
                self.wall_skill.use()   # 激活穿墙
            elif skill_type == "length":
                self.length_skill.use() # 激活变长
                for _ in range(3):
                    self.snake.append(self.snake[-1])
                self.score += 3
            self.skill_food = None  # 清除技能食物
        # 没吃到食物则移除尾部
        else:
            self.snake.pop()

    def draw_game(self):
        """绘制游戏画面"""
        screen.fill(BLACK)
        
        # 绘制蛇
        for i, segment in enumerate(self.snake):
            color = GREEN if i == 0 else (0, 200, 0)  # 蛇头绿色更亮
            pygame.draw.rect(screen, color, (segment[0], segment[1], BLOCK_SIZE, BLOCK_SIZE))
            # 绘制蛇的边框
            pygame.draw.rect(screen, WHITE, (segment[0], segment[1], BLOCK_SIZE, BLOCK_SIZE), 1)
        
        # 绘制普通食物
        pygame.draw.rect(screen, RED, (self.food[0], self.food[1], BLOCK_SIZE, BLOCK_SIZE))
        
        # 绘制技能食物
        if self.skill_food:
            pos, skill_type = self.skill_food
            if skill_type == "speed":
                color = YELLOW
            elif skill_type == "wall":
                color = PURPLE
            else:
                color = CYAN
            pygame.draw.rect(screen, color, (pos[0], pos[1], BLOCK_SIZE, BLOCK_SIZE))
        
        # 绘制分数和技能信息
        score_text = FONT.render(f"分数: {self.score}", True, WHITE)
        screen.blit(score_text, (10, 10))
        
        # 技能状态显示
        speed_text = FONT.render(f"加速(1): {'冷却中' if self.speed_skill.current_cooldown > 0 else '可用'}", True, YELLOW)
        wall_text = FONT.render(f"穿墙(2): {'冷却中' if self.wall_skill.current_cooldown > 0 else '可用'}", True, PURPLE)
        length_text = FONT.render(f"变长(3): {'冷却中' if self.length_skill.current_cooldown > 0 else '可用'}", True, CYAN)
        screen.blit(speed_text, (10, 50))
        screen.blit(wall_text, (10, 90))
        screen.blit(length_text, (10, 130))
        
        # 游戏结束提示
        if self.game_over:
            game_over_text = FONT.render("游戏结束！按R重新开始", True, RED)
            screen.blit(game_over_text, (WIDTH//2 - 180, HEIGHT//2))
        
        # 暂停提示
        if self.paused and not self.game_over:
            pause_text = FONT.render("暂停中！按P继续", True, YELLOW)
            screen.blit(pause_text, (WIDTH//2 - 120, HEIGHT//2))
        
        pygame.display.update()

    def run(self):
        """游戏主循环"""
        while True:
            self.handle_events()
            self.update_game()
            self.draw_game()
            
            # 动态调整帧率（加速技能生效时帧率翻倍）
            current_fps = FPS * 2 if self.speed_skill.active else FPS
            clock.tick(current_fps)

# 启动游戏
if __name__ == "__main__":
    game = SkillSnakeGame()
    game.run()