import random
import pgzrun

# ====================== 全局基础常量 ======================
WIDTH = 600
HEIGHT = 600
GRID = 20
GRID_NUM_X = WIDTH // GRID
GRID_NUM_Y = HEIGHT // GRID
FORBID_RANGE = 10

# 难度配置：Lv1=40，Lv2=30，Lv3=20
LEVEL_CFG = {
    1: (8, 40),
    2: (4, 30),
    3: (2, 20)
}
DEFAULT_LEVEL = 2

# 色块颜色
BG_COLOR = (0, 0, 0)
OBSTACLE_COLOR = (139, 69, 19)
FOOD_FALLBACK_COLOR = (255, 0, 0)
SNAKE_BODY_COLOR = (0, 0, 255)
SNAKE_HEAD_COLOR = (0, 255, 0)

# ====================== 全局状态 ======================
game_state = "menu"
current_level = DEFAULT_LEVEL
move_frame_gap = LEVEL_CFG[current_level][0]
win_length = LEVEL_CFG[current_level][1]

snake = []
dir_x, dir_y = 1, 0
next_dir_x, next_dir_y = 1, 0
frame_count = 0

score = 0
food_eaten_count = 0
food_pos = (0, 0)
food_actor = None
obstacles = []
has_revived = False
bgm_obj = None       # 背景音乐对象
bgm_playing = False  # BGM播放标记

# 全部音效变量（bgm并入sounds）
sound_start = None
sound_eat = None
sound_revive = None
sound_gameover = None
sound_win = None

# ====================== 素材加载 ======================
def load_resources():
    global food_actor, bgm_obj, sound_start, sound_eat, sound_revive, sound_gameover, sound_win
    # 食物图片
    try:
        food_actor = Actor("food")
    except:
        food_actor = None
    # 所有音效统一从sounds加载
    try:
        bgm_obj = sounds.bgm
    except Exception as e:
        print("bgm.wav 加载失败：", e)
        bgm_obj = None
    try:
        sound_start = sounds.start
    except:
        sound_start = None
    try:
        sound_eat = sounds.eat
    except:
        sound_eat = None
    try:
        sound_revive = sounds.revive
    except:
        sound_revive = None
    try:
        sound_gameover = sounds.game_over
    except:
        sound_gameover = None
    try:
        sound_win = sounds.win
    except:
        sound_win = None

# ====================== 统一停止背景音乐函数（修复重叠核心） ======================
def stop_bgm():
    global bgm_playing
    if bgm_obj is not None:
        bgm_obj.stop()
    bgm_playing = False

# ====================== 禁区判断 ======================
def is_in_head_forbid_area(pos):
    x, y = pos
    hx, hy = snake[0]
    if dir_x == 1:
        if x > hx and (x - hx) <= FORBID_RANGE and abs(y - hy) <= 1:
            return True
    elif dir_x == -1:
        if x < hx and (hx - x) <= FORBID_RANGE and abs(y - hy) <= 1:
            return True
    elif dir_y == 1:
        if y > hy and (y - hy) <= FORBID_RANGE and abs(x - hx) <= 1:
            return True
    elif dir_y == -1:
        if y < hy and (hy - y) <= FORBID_RANGE and abs(x - hx) <= 1:
            return True
    return False

# ====================== 游戏工具函数 ======================
def reset_game():
    global snake, dir_x, dir_y, next_dir_x, next_dir_y
    global score, food_eaten_count, obstacles, frame_count, has_revived
    center_x = GRID_NUM_X // 2
    center_y = GRID_NUM_Y // 2
    snake = [(center_x, center_y), (center_x - 1, center_y), (center_x - 2, center_y)]
    dir_x, dir_y = 1, 0
    next_dir_x, next_dir_y = 1, 0
    frame_count = 0
    score = 0
    food_eaten_count = 0
    obstacles.clear()
    has_revived = False
    spawn_food()

def revive_player():
    global snake, dir_x, dir_y, next_dir_x, next_dir_y, frame_count, obstacles, has_revived
    if sound_revive:
        sound_revive.play()
    center_x = GRID_NUM_X // 2
    center_y = GRID_NUM_Y // 2
    length = len(snake)
    snake = []
    for i in range(length):
        snake.append((center_x - i, center_y))
    dir_x, dir_y = 1, 0
    next_dir_x, next_dir_y = 1, 0
    frame_count = 0
    obstacles.clear()
    has_revived = True
    spawn_food()

def get_empty_positions():
    empty = []
    snake_set = set(snake)
    obs_set = set(obstacles)
    for x in range(GRID_NUM_X):
        for y in range(GRID_NUM_Y):
            pos = (x, y)
            if pos not in snake_set and pos not in obs_set and not is_in_head_forbid_area(pos):
                empty.append(pos)
    return empty

def spawn_food():
    global food_pos
    empty = get_empty_positions()
    food_pos = random.choice(empty)

def spawn_obstacle_group():
    empty = get_empty_positions()
    if len(empty) < 3:
        return
    start = random.choice(empty)
    sx, sy = start
    shapes = [
        [(sx, sy), (sx+1, sy), (sx+2, sy)],
        [(sx, sy), (sx, sy+1), (sx, sy+2)],
        [(sx, sy), (sx+1, sy), (sx, sy+1)]
    ]
    shape = random.choice(shapes)
    fixed_shape = []
    for (x, y) in shape:
        fx = x % GRID_NUM_X
        fy = y % GRID_NUM_Y
        fixed_shape.append((fx, fy))
    snake_set = set(snake)
    obs_set = set(obstacles)
    food_set = {food_pos}
    valid = True
    for p in fixed_shape:
        if p in snake_set or p in obs_set or p in food_set or is_in_head_forbid_area(p):
            valid = False
            break
    if valid:
        obstacles.extend(fixed_shape)

def check_win():
    return len(snake) >= win_length

# ====================== 游戏更新 ======================
def update():
    global game_state, frame_count, dir_x, dir_y, next_dir_x, next_dir_y
    global score, food_eaten_count, snake, obstacles, bgm_playing

    # 进入游戏界面且未播放BGM时启动，play(-1) = 无限循环
    if game_state == "play" and not bgm_playing and bgm_obj is not None:
        bgm_obj.play(-1)
        bgm_playing = True
        print("BGM启动成功")

    # 非对局界面跳过逻辑
    if game_state == "menu" or game_state in ("death_revive", "gameover", "win"):
        return

    frame_count += 1

    # 禁止反向掉头
    if not (next_dir_x == -dir_x and next_dir_y == -dir_y):
        dir_x, dir_y = next_dir_x, next_dir_y

    if frame_count % move_frame_gap != 0:
        return

    head_x, head_y = snake[0]
    new_x = (head_x + dir_x) % GRID_NUM_X
    new_y = (head_y + dir_y) % GRID_NUM_Y
    new_head = (new_x, new_y)

    # 撞障碍/自身死亡
    if new_head in snake or new_head in obstacles:
        if sound_gameover:
            sound_gameover.play()
        if not has_revived:
            game_state = "death_revive"
        else:
            game_state = "gameover"
        return

    # 吃到食物
    if new_head == food_pos:
        snake.insert(0, new_head)
        score += 1
        food_eaten_count += 1
        if food_eaten_count % 2 == 0:
            spawn_obstacle_group()
        spawn_food()
        if sound_eat:
            sound_eat.play()
        if check_win():
            game_state = "win"
            if sound_win:
                sound_win.play()
        return
    snake.insert(0, new_head)
    snake.pop()

# ====================== 按键控制 ======================
def on_key_down(key):
    global game_state, current_level, move_frame_gap, win_length
    global next_dir_x, next_dir_y

    # ESC退出：先停止BGM再关闭程序
    if key == keys.ESCAPE:
        stop_bgm()
        exit()

    # 难度菜单选关：进入游戏前停止旧BGM
    if game_state == "menu":
        if key == keys.K_1:
            current_level = 1
        elif key == keys.K_2:
            current_level = 2
        elif key == keys.K_3:
            current_level = 3
        else:
            return
        move_frame_gap, win_length = LEVEL_CFG[current_level]
        stop_bgm()
        reset_game()
        if sound_start:
            sound_start.play()
        game_state = "play"
        return

    # 死亡复活界面
    if game_state == "death_revive":
        if key == keys.Q:
            revive_player()
            game_state = "play"
        elif key == keys.R:
            # 按R重开：先停止BGM再重置
            stop_bgm()
            reset_game()
            if sound_start:
                sound_start.play()
            game_state = "play"
        return

    # 胜利/彻底失败界面按R重开
    if game_state in ("gameover", "win"):
        if key == keys.R:
            # 按R重开：先停止BGM再重置
            stop_bgm()
            reset_game()
            if sound_start:
                sound_start.play()
            game_state = "play"
        return

    # WASD移动
    if game_state == "play":
        if key == keys.W:
            next_dir_x, next_dir_y = 0, -1
        elif key == keys.S:
            next_dir_x, next_dir_y = 0, 1
        elif key == keys.A:
            next_dir_x, next_dir_y = -1, 0
        elif key == keys.D:
            next_dir_x, next_dir_y = 1, 0

# ====================== 绘制 ======================
def draw():
    screen.fill(BG_COLOR)

    # 难度菜单文字 40 / 30 / 20
    if game_state == "menu":
        screen.draw.text("SNAKE GAME", center=(WIDTH//2, 120), fontsize=48, color="white")
        screen.draw.text("Press 1 / 2 / 3 to select level", center=(WIDTH//2, 180), fontsize=24, color="gray")
        screen.draw.text("Level 1: Slow | Target Length 40", center=(WIDTH//2, 230), fontsize=20, color="green")
        screen.draw.text("Level 2: Normal | Target Length 30 (Default)", center=(WIDTH//2, 270), fontsize=20, color="blue")
        screen.draw.text("Level 3: Fast | Target Length 20", center=(WIDTH//2, 310), fontsize=20, color="red")
        screen.draw.text("WASD Move | ESC Quit", center=(WIDTH//2, 400), fontsize=22, color="lightgray")
        return

    # 绘制障碍物
    for (x, y) in obstacles:
        rect = Rect(x * GRID, y * GRID, GRID, GRID)
        screen.draw.filled_rect(rect, OBSTACLE_COLOR)

    # 绘制食物
    fx, fy = food_pos
    food_center_x = fx * GRID + GRID / 2
    food_center_y = fy * GRID + GRID / 2
    if food_actor:
        food_actor.pos = (food_center_x, food_center_y)
        food_actor.draw()
    else:
        food_rect = Rect(fx * GRID, fy * GRID, GRID, GRID)
        screen.draw.filled_rect(food_rect, FOOD_FALLBACK_COLOR)

    # 绘制蛇身体
    for idx in range(1, len(snake)):
        x, y = snake[idx]
        body_rect = Rect(x * GRID, y * GRID, GRID, GRID)
        screen.draw.filled_rect(body_rect, SNAKE_BODY_COLOR)
    # 蛇头
    hx, hy = snake[0]
    head_rect = Rect(hx * GRID, hy * GRID, GRID, GRID)
    screen.draw.filled_rect(head_rect, SNAKE_HEAD_COLOR)

    # 左上角文字信息
    screen.draw.text(f"Score: {score}", (10, 10), fontsize=20, color="white")
    screen.draw.text(f"Level: {current_level}", (10, 35), fontsize=20, color="white")
    screen.draw.text(f"Obstacle Blocks: {len(obstacles)}", (10, 60), fontsize=20, color="lightgray")

    # 死亡复活遮罩
    if game_state == "death_revive":
        screen.draw.filled_rect(Rect(0, 0, WIDTH, HEIGHT), (0, 0, 0, 180))
        screen.draw.text("YOU DIED", center=(WIDTH//2, HEIGHT//2 - 60), fontsize=60, color="orange")
        screen.draw.text("Press Q to Revive (Keep score & length)", center=(WIDTH//2, HEIGHT//2 - 10), fontsize=22, color="white")
        screen.draw.text("Press R to Restart Full Game", center=(WIDTH//2, HEIGHT//2 + 30), fontsize=22, color="white")
        screen.draw.text("ESC to Quit", center=(WIDTH//2, HEIGHT//2 + 70), fontsize=20, color="lightgray")

    # Game Over遮罩
    if game_state == "gameover":
        screen.draw.filled_rect(Rect(0, 0, WIDTH, HEIGHT), (0, 0, 0, 180))
        screen.draw.text("GAME OVER", center=(WIDTH//2, HEIGHT//2 - 40), fontsize=60, color="red")
        screen.draw.text("R Restart | ESC Quit", center=(WIDTH//2, HEIGHT//2 + 30), fontsize=24, color="white")

    # 胜利遮罩
    if game_state == "win":
        screen.draw.filled_rect(Rect(0, 0, WIDTH, HEIGHT), (0, 0, 0, 180))
        screen.draw.text("YOU WIN!", center=(WIDTH//2, HEIGHT//2 - 50), fontsize=60, color="green")
        screen.draw.text("All space filled completely!", center=(WIDTH//2, HEIGHT//2 + 10), fontsize=22, color="white")
        screen.draw.text("R Restart | ESC Quit", center=(WIDTH//2, HEIGHT//2 + 50), fontsize=24, color="white")

# 加载素材
load_resources()
# 启动游戏
pgzrun.go()