import pygame
import random
import sys
import os

# -------------------------- 基础配置 --------------------------
WIDTH = 600
HEIGHT = 600
BLOCK_SIZE = 20

# 难度对应速度（数字越大蛇移动越快）
LEVEL_SPEEDS = {
    1: 10,    # Level1 最慢
    2: 15,   # Level2 正常速度
    3: 22    # Level3 较快速度
}

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)
GREEN = (0, 255, 0)
BLUE = (0, 0, 255)
GRAY = (180, 180, 180)
OBSTACLE_COLOR = (139, 69, 19)  # 马鞍棕障碍物

# 场地总格子数量，用于胜利判定
TOTAL_BLOCKS = (WIDTH // BLOCK_SIZE) * (HEIGHT // BLOCK_SIZE)

# -------------------------- 资源加载（带容错） --------------------------
assets_dir = os.path.join(os.path.dirname(__file__), "assets")

# 四方向蛇头图片
snake_head_up = None
snake_head_down = None
snake_head_left = None
snake_head_right = None

# 蛇身、食物、障碍物图片
snake_body_img = None
food_img = None
obstacle_img = None

# 音效
eat_sound = None
game_over_sound = None
win_sound = None

def load_assets():
    """加载所有素材，加载失败不影响游戏运行"""
    global snake_head_up, snake_head_down, snake_head_left, snake_head_right
    global snake_body_img, food_img, obstacle_img, eat_sound, game_over_sound, win_sound

    # 加载图片资源
    try:
        # 四方向蛇头
        up_path = os.path.join(assets_dir, "snake_head_up.png")
        if os.path.exists(up_path):
            snake_head_up = pygame.image.load(up_path).convert_alpha()
            snake_head_up = pygame.transform.scale(snake_head_up, (BLOCK_SIZE, BLOCK_SIZE))
        
        down_path = os.path.join(assets_dir, "snake_head_down.png")
        if os.path.exists(down_path):
            snake_head_down = pygame.image.load(down_path).convert_alpha()
            snake_head_down = pygame.transform.scale(snake_head_down, (BLOCK_SIZE, BLOCK_SIZE))
        
        left_path = os.path.join(assets_dir, "snake_head_left.png")
        if os.path.exists(left_path):
            snake_head_left = pygame.image.load(left_path).convert_alpha()
            snake_head_left = pygame.transform.scale(snake_head_left, (BLOCK_SIZE, BLOCK_SIZE))
        
        right_path = os.path.join(assets_dir, "snake_head_right.png")
        if os.path.exists(right_path):
            snake_head_right = pygame.image.load(right_path).convert_alpha()
            snake_head_right = pygame.transform.scale(snake_head_right, (BLOCK_SIZE, BLOCK_SIZE))

        # 蛇身
        body_path = os.path.join(assets_dir, "snake_body.png")
        if os.path.exists(body_path):
            snake_body_img = pygame.image.load(body_path).convert_alpha()
            snake_body_img = pygame.transform.scale(snake_body_img, (BLOCK_SIZE, BLOCK_SIZE))
        
        # 食物
        food_path = os.path.join(assets_dir, "food.png")
        if os.path.exists(food_path):
            food_img = pygame.image.load(food_path).convert_alpha()
            food_img = pygame.transform.scale(food_img, (BLOCK_SIZE, BLOCK_SIZE))
        
        # 障碍物图片
        obs_path = os.path.join(assets_dir, "obstacle.png")
        if os.path.exists(obs_path):
            obstacle_img = pygame.image.load(obs_path).convert_alpha()
            obstacle_img = pygame.transform.scale(obstacle_img, (BLOCK_SIZE, BLOCK_SIZE))
    except:
        pass

    # 加载音效资源
    try:
        pygame.mixer.init()
        eat_path = os.path.join(assets_dir, "eat.wav")
        if os.path.exists(eat_path):
            eat_sound = pygame.mixer.Sound(eat_path)
        
        over_path = os.path.join(assets_dir, "game_over.wav")
        if os.path.exists(over_path):
            game_over_sound = pygame.mixer.Sound(over_path)
            
        win_path = os.path.join(assets_dir, "win.wav")
        if os.path.exists(win_path):
            win_sound = pygame.mixer.Sound(win_path)
    except:
        pass

# -------------------------- 游戏核心类 --------------------------
class SnakeGame:
    def __init__(self):
        pygame.init()
        # 关闭系统自带键盘长按延迟（关键优化1）
        pygame.key.set_repeat(0, 0)
        self.screen = pygame.display.set_mode((WIDTH, HEIGHT))
        pygame.display.set_caption("贪吃蛇 - WASD控制")
        self.clock = pygame.time.Clock()
        
        # 多级字体：标题、大标题、普通文本
        self.title_font = pygame.font.Font(None, 64)
        self.big_font = pygame.font.Font(None, 48)
        self.menu_font = pygame.font.Font(None, 32)
        self.font = pygame.font.Font(None, 25)
        
        # 加载素材
        load_assets()
        
        # 游戏状态：menu(菜单选难度) / playing(游戏中) / game_over(结束) / win(胜利)
        self.state = 'menu'
        # 默认难度为Level2
        self.current_level = 2
        self.speed = LEVEL_SPEEDS[self.current_level]
        
        # 障碍物列表、吃食物计数（每2个生成一组三连障碍）
        self.obstacles = []
        self.eat_count_for_obs = 0
        
        # 初始化游戏数据（菜单状态下不运行逻辑）
        self.reset_game()

    def reset_game(self):
        """重置游戏状态，保留当前难度，清空障碍物"""
        self.snake = [
            (WIDTH // 2, HEIGHT // 2),
            (WIDTH // 2 - BLOCK_SIZE, HEIGHT // 2),
            (WIDTH // 2 - 2 * BLOCK_SIZE, HEIGHT // 2)
        ]
        # 初始方向向右
        self.dx = BLOCK_SIZE
        self.dy = 0
        self.obstacles = []
        self.eat_count_for_obs = 0
        self.spawn_food()
        self.score = 0

    def get_all_occupied_pos(self):
        """获取所有被占用的格子：蛇身 + 障碍物"""
        occupied = set(self.snake)
        occupied.update(set(self.obstacles))
        return occupied

    def spawn_food(self):
        """随机生成食物，避开蛇身和障碍物；无空位则触发胜利"""
        occupied = self.get_all_occupied_pos()
        # 所有格子都被占满 → 游戏胜利
        if len(occupied) >= TOTAL_BLOCKS:
            self.state = "win"
            if win_sound:
                win_sound.play()
            return
        
        # 正常刷新食物
        while True:
            x = random.randrange(0, WIDTH - BLOCK_SIZE, BLOCK_SIZE)
            y = random.randrange(0, HEIGHT - BLOCK_SIZE, BLOCK_SIZE)
            pos = (x, y)
            if pos not in occupied:
                self.food = pos
                break

    def spawn_triple_obstacle(self):
        """生成一组3块相连障碍物，随机形状：横条/竖条/L型"""
        occupied = self.get_all_occupied_pos()
        occupied.add(self.food)
        grid_w = WIDTH // BLOCK_SIZE
        grid_h = HEIGHT // BLOCK_SIZE

        # 三种三连障碍物形状偏移量
        shapes = [
            [(0,0), (BLOCK_SIZE,0), (BLOCK_SIZE*2,0)],    # 横向直线
            [(0,0), (0,BLOCK_SIZE), (0,BLOCK_SIZE*2)],    # 纵向直线
            [(0,0), (BLOCK_SIZE,0), (0,BLOCK_SIZE)]       # L型拐角
        ]
        shape = random.choice(shapes)

        # 循环找合法生成点
        while True:
            # 随机基础坐标（保证不贴边界）
            base_x = random.randrange(BLOCK_SIZE, WIDTH - BLOCK_SIZE*2, BLOCK_SIZE)
            base_y = random.randrange(BLOCK_SIZE, HEIGHT - BLOCK_SIZE*2, BLOCK_SIZE)
            block_group = []
            conflict = False

            for dx, dy in shape:
                px = base_x + dx
                py = base_y + dy
                pos = (px, py)
                if pos in occupied:
                    conflict = True
                    break
                block_group.append(pos)
            
            if not conflict:
                # 无冲突，全部加入障碍物列表
                self.obstacles.extend(block_group)
                break

    def check_wall_collision(self, head_x, head_y):
        """穿墙逻辑：边界穿越"""
        if head_x < 0:
            head_x = WIDTH - BLOCK_SIZE
        elif head_x >= WIDTH:
            head_x = 0
        
        if head_y < 0:
            head_y = HEIGHT - BLOCK_SIZE
        elif head_y >= HEIGHT:
            head_y = 0
        
        return head_x, head_y

    def check_collision(self, head):
        """碰撞检测：自身、障碍物"""
        # 撞到自己
        if head in self.snake[1:]:
            return True
        # 撞到障碍物
        if head in self.obstacles:
            return True
        return False

    def handle_input(self):
        """修复无延迟输入：分事件+实时按键采样"""
        event_list = pygame.event.get()
        keys = pygame.key.get_pressed() # 实时读取所有按住的按键（消除长按延迟）

        # 全局退出事件
        for event in event_list:
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            # 一次性按键触发（菜单/重试用）
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    pygame.quit()
                    sys.exit()
                # 菜单选难度
                if self.state == 'menu':
                    if event.key == pygame.K_1:
                        self.current_level = 1
                        self.speed = LEVEL_SPEEDS[1]
                        self.reset_game()
                        self.state = 'playing'
                    elif event.key == pygame.K_2:
                        self.current_level = 2
                        self.speed = LEVEL_SPEEDS[2]
                        self.reset_game()
                        self.state = 'playing'
                    elif event.key == pygame.K_3:
                        self.current_level = 3
                        self.speed = LEVEL_SPEEDS[3]
                        self.reset_game()
                        self.state = 'playing'
                # 结束/胜利界面重试
                elif self.state in ["game_over", "win"]:
                    if event.key == pygame.K_r:
                        self.reset_game()
                        self.state = 'playing'

        # 游戏中实时WASD控制（核心修复：实时采样无延迟）
        if self.state == "playing":
            # W上，禁止反向向下
            if keys[pygame.K_w] and self.dy != BLOCK_SIZE:
                self.dy = -BLOCK_SIZE
                self.dx = 0
            # S下，禁止反向向上
            elif keys[pygame.K_s] and self.dy != -BLOCK_SIZE:
                self.dy = BLOCK_SIZE
                self.dx = 0
            # A左，禁止反向向右
            elif keys[pygame.K_a] and self.dx != BLOCK_SIZE:
                self.dx = -BLOCK_SIZE
                self.dy = 0
            # D右，禁止反向向左
            elif keys[pygame.K_d] and self.dx != -BLOCK_SIZE:
                self.dx = BLOCK_SIZE
                self.dy = 0

    def update(self):
        """更新游戏逻辑，仅游戏中运行"""
        if self.state != 'playing':
            return

        head_x, head_y = self.snake[0]
        new_head = (head_x + self.dx, head_y + self.dy)
        new_head = self.check_wall_collision(*new_head)
        
        # 撞到自身/障碍物 → Game Over
        if self.check_collision(new_head):
            self.state = 'game_over'
            if game_over_sound:
                game_over_sound.play()
            return
        
        self.snake.insert(0, new_head)
        
        # 吃到食物逻辑
        if new_head == self.food:
            self.score += 1
            self.eat_count_for_obs += 1
            if eat_sound:
                eat_sound.play()
            # 每吃2个食物，生成一组3连块障碍物
            if self.eat_count_for_obs >= 2:
                self.spawn_triple_obstacle()
                self.eat_count_for_obs = 0
            self.spawn_food()
        # 没吃到食物：删除尾部，保持长度
        else:
            self.snake.pop()

    def draw_snake_head(self, pos):
        """根据当前方向绘制对应蛇头"""
        x, y = pos
        # 向右
        if self.dx == BLOCK_SIZE:
            if snake_head_right:
                self.screen.blit(snake_head_right, (x, y))
            else:
                pygame.draw.rect(self.screen, GREEN, (x, y, BLOCK_SIZE, BLOCK_SIZE))
        # 向左
        elif self.dx == -BLOCK_SIZE:
            if snake_head_left:
                self.screen.blit(snake_head_left, (x, y))
            else:
                pygame.draw.rect(self.screen, GREEN, (x, y, BLOCK_SIZE, BLOCK_SIZE))
        # 向上
        elif self.dy == -BLOCK_SIZE:
            if snake_head_up:
                self.screen.blit(snake_head_up, (x, y))
            else:
                pygame.draw.rect(self.screen, GREEN, (x, y, BLOCK_SIZE, BLOCK_SIZE))
        # 向下
        elif self.dy == BLOCK_SIZE:
            if snake_head_down:
                self.screen.blit(snake_head_down, (x, y))
            else:
                pygame.draw.rect(self.screen, GREEN, (x, y, BLOCK_SIZE, BLOCK_SIZE))

    def draw_menu(self):
        """绘制难度选择菜单"""
        # 标题
        title_text = self.title_font.render("SNAKE GAME", True, GREEN)
        self.screen.blit(title_text, (WIDTH//2 - title_text.get_width()//2, 120))
        
        # 副标题
        sub_text = self.menu_font.render("Select Difficulty", True, WHITE)
        self.screen.blit(sub_text, (WIDTH//2 - sub_text.get_width()//2, 200))

        # 三个难度选项
        level1_text = self.menu_font.render("Press 1 : Level 1 (Slow)", True, GRAY)
        self.screen.blit(level1_text, (WIDTH//2 - level1_text.get_width()//2, 280))
        
        level2_text = self.menu_font.render("Press 2 : Level 2 (Normal)", True, WHITE)
        self.screen.blit(level2_text, (WIDTH//2 - level2_text.get_width()//2, 320))
        
        level3_text = self.menu_font.render("Press 3 : Level 3 (Fast)", True, GRAY)
        self.screen.blit(level3_text, (WIDTH//2 - level3_text.get_width()//2, 360))

        # 底部规则提示
        tip_text = self.font.render("WASD move | Every 2 food spawn 3 connected obstacles | ESC exit", True, GRAY)
        self.screen.blit(tip_text, (WIDTH//2 - tip_text.get_width()//2, 480))

    def draw(self):
        """绘制全部画面，分状态渲染"""
        self.screen.fill(BLACK)

        # 菜单界面
        if self.state == 'menu':
            self.draw_menu()
        
        # 游戏界面（游戏中+结束+胜利都基于游戏画面）
        else:
            # 绘制障碍物
            for obs_pos in self.obstacles:
                if obstacle_img:
                    self.screen.blit(obstacle_img, obs_pos)
                else:
                    pygame.draw.rect(self.screen, OBSTACLE_COLOR, (obs_pos[0], obs_pos[1], BLOCK_SIZE, BLOCK_SIZE))

            # 绘制食物
            if food_img:
                self.screen.blit(food_img, self.food)
            else:
                pygame.draw.rect(self.screen, RED, (self.food[0], self.food[1], BLOCK_SIZE, BLOCK_SIZE))

            # 绘制蛇身和蛇头
            for i, block in enumerate(self.snake):
                if i == 0:
                    self.draw_snake_head(block)
                else:
                    if snake_body_img:
                        self.screen.blit(snake_body_img, block)
                    else:
                        pygame.draw.rect(self.screen, BLUE, (block[0], block[1], BLOCK_SIZE, BLOCK_SIZE))

            # 左上角信息
            score_text = self.font.render(f"Score: {self.score}", True, WHITE)
            self.screen.blit(score_text, (10, 10))
            
            level_text = self.font.render(f"Level: {self.current_level}", True, GRAY)
            self.screen.blit(level_text, (10, 35))
            
            obs_count_text = self.font.render(f"Obstacle Blocks: {len(self.obstacles)}", True, GRAY)
            self.screen.blit(obs_count_text, (10, 60))

            # 游戏结束遮罩界面
            if self.state == 'game_over':
                overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
                overlay.fill((0, 0, 0, 180))
                self.screen.blit(overlay, (0, 0))

                over_title = self.big_font.render("GAME OVER", True, RED)
                self.screen.blit(over_title, (WIDTH//2 - over_title.get_width()//2, HEIGHT//2 - 70))
                
                retry_text = self.menu_font.render("R - Retry", True, WHITE)
                self.screen.blit(retry_text, (WIDTH//2 - retry_text.get_width()//2, HEIGHT//2 + 10))
                
                close_text = self.menu_font.render("ESC - Close", True, WHITE)
                self.screen.blit(close_text, (WIDTH//2 - close_text.get_width()//2, HEIGHT//2 + 50))
            
            # 游戏胜利遮罩界面
            elif self.state == 'win':
                overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
                overlay.fill((0, 0, 0, 180))
                self.screen.blit(overlay, (0, 0))

                win_title = self.big_font.render("YOU WIN!", True, GREEN)
                self.screen.blit(win_title, (WIDTH//2 - win_title.get_width()//2, HEIGHT//2 - 70))
                
                win_tip = self.menu_font.render("All space filled completely!", True, WHITE)
                self.screen.blit(win_tip, (WIDTH//2 - win_tip.get_width()//2, HEIGHT//2 - 20))
                
                retry_text = self.menu_font.render("R - Retry", True, WHITE)
                self.screen.blit(retry_text, (WIDTH//2 - retry_text.get_width()//2, HEIGHT//2 + 20))
                
                close_text = self.menu_font.render("ESC - Close", True, WHITE)
                self.screen.blit(close_text, (WIDTH//2 - close_text.get_width()//2, HEIGHT//2 + 60))

        pygame.display.flip()

    def run(self):
        """游戏主循环"""
        while True:
            self.handle_input()
            self.update()
            self.draw()
            self.clock.tick(self.speed)

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