import pygame
import math
import sys

# ==================== 初始化 ====================
pygame.init()

# 窗口设置
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 塔防 · 带界面")
clock = pygame.time.Clock()
FPS = 60

# 颜色定义
COLOR_BG = (30, 47, 58)
COLOR_GRID = (46, 74, 92)
COLOR_PATH = (79, 111, 127)
COLOR_PATH_LIGHT = (106, 141, 156)
COLOR_BLOCKED = (120, 60, 60)
COLOR_TOWER_BASE = (62, 107, 126)
COLOR_TOWER_TOP = (142, 201, 224)
COLOR_TOWER_GUN = (255, 228, 132)
COLOR_ENEMY = (184, 74, 74)
COLOR_BULLET = (255, 224, 102)
COLOR_TEXT = (191, 217, 232)
COLOR_GOLD = (255, 217, 102)
COLOR_HP_BG = (30, 30, 30)
COLOR_HP_GREEN = (75, 224, 75)
COLOR_HP_YELLOW = (249, 215, 28)
COLOR_HP_RED = (240, 78, 78)

# 字体 —— 使用 pygame 内置默认字体，避免 SysFont 扫描系统字体出错
FONT = pygame.font.Font(None, 24)
FONT_SMALL = pygame.font.Font(None, 20)
FONT_BIG = pygame.font.Font(None, 48)

# ==================== 网格配置 ====================
GRID_SIZE = 40
COLS = WIDTH // GRID_SIZE       # 22
ROWS = HEIGHT // GRID_SIZE      # 15

# ==================== 路径定义 ====================
waypoints = [
    (0, 80),
    (360, 80),
    (360, 240),
    (640, 240),
    (640, 400),
    (880, 400),
    (920, 400),
]

# 标记不可建造网格
path_grid = set()


def distance_to_segment(px, py, x1, y1, x2, y2):
    """点到线段的最短距离"""
    vx, vy = x2 - x1, y2 - y1
    dx, dy = px - x1, py - y1
    dot = dx * vx + dy * vy
    len2 = vx * vx + vy * vy
    if len2 == 0:
        return math.hypot(px - x1, py - y1)
    t = max(0, min(1, dot / len2))
    proj_x = x1 + t * vx
    proj_y = y1 + t * vy
    return math.hypot(px - proj_x, py - proj_y)


def mark_path_grid():
    """把路径附近的格子标记为不可建造"""
    for i in range(len(waypoints) - 1):
        x1, y1 = waypoints[i]
        x2, y2 = waypoints[i + 1]
        for col in range(COLS):
            for row in range(ROWS):
                cx = col * GRID_SIZE + GRID_SIZE / 2
                cy = row * GRID_SIZE + GRID_SIZE / 2
                if distance_to_segment(cx, cy, x1, y1, x2, y2) < 28:
                    path_grid.add((col, row))


mark_path_grid()

# ==================== 游戏状态 ====================
lives = 10
gold = 250
wave = 1

towers = []
enemies = []
bullets = []

enemies_this_wave = 0
enemies_spawned = 0
spawn_timer = 0
SPAWN_DELAY = 45

wave_in_progress = False
wave_cooldown = 0

# 塔属性
TOWER_COST = 50
TOWER_RANGE = 110
TOWER_COOLDOWN_MAX = 30
TOWER_DAMAGE = 1

# 敌人属性
ENEMY_SPEED = 1.2
ENEMY_REWARD = 30
ENEMY_BASE_HP = 2

# 鼠标网格
mouse_col = -1
mouse_row = -1

game_over = False


# ==================== 游戏逻辑函数 ====================
def start_wave():
    global wave_in_progress, enemies_this_wave, enemies_spawned, spawn_timer, gold
    if wave_in_progress:
        return
    wave_in_progress = True
    enemies_this_wave = 5 + wave * 2
    enemies_spawned = 0
    spawn_timer = 0
    gold += 100
    print(f"[波次] 第 {wave} 波开始! 敌人数量: {enemies_this_wave}")


def spawn_enemy():
    global enemies_spawned
    if enemies_spawned >= enemies_this_wave:
        return
    start_x, start_y = waypoints[0]
    offset = (pygame.time.get_ticks() % 16) - 8
    hp = ENEMY_BASE_HP + int(wave * 0.8)
    enemies.append({
        "x": float(start_x),
        "y": float(start_y + offset),
        "hp": hp,
        "max_hp": hp,
        "speed": ENEMY_SPEED,
        "wp_index": 1,
        "size": 14,
        "reward": ENEMY_REWARD,
    })
    enemies_spawned += 1


def update_enemies():
    global lives, game_over
    for i in range(len(enemies) - 1, -1, -1):
        e = enemies[i]
        if e["wp_index"] >= len(waypoints):
            lives -= 1
            enemies.pop(i)
            if lives <= 0:
                lives = 0
                game_over = True
            continue

        tx, ty = waypoints[e["wp_index"]]
        dx, dy = tx - e["x"], ty - e["y"]
        dist = math.hypot(dx, dy)

        if dist < 2:
            e["wp_index"] += 1
            if e["wp_index"] >= len(waypoints):
                lives -= 1
                enemies.pop(i)
                if lives <= 0:
                    lives = 0
                    game_over = True
            continue

        e["x"] += (dx / dist) * e["speed"]
        e["y"] += (dy / dist) * e["speed"]


def update_towers():
    for tower in towers:
        if tower["cooldown"] > 0:
            tower["cooldown"] -= 1
            continue

        closest = None
        min_dist = float("inf")
        for e in enemies:
            d = math.hypot(e["x"] - tower["x"], e["y"] - tower["y"])
            if d <= TOWER_RANGE and d < min_dist:
                min_dist = d
                closest = e

        if closest is not None:
            tower["cooldown"] = TOWER_COOLDOWN_MAX
            bullets.append({
                "x": tower["x"],
                "y": tower["y"],
                "target": closest,
                "speed": 6,
                "damage": TOWER_DAMAGE,
            })


def update_bullets():
    global gold
    for i in range(len(bullets) - 1, -1, -1):
        b = bullets[i]
        target = b["target"]

        if target not in enemies:
            bullets.pop(i)
            continue

        dx, dy = target["x"] - b["x"], target["y"] - b["y"]
        dist = math.hypot(dx, dy)

        if dist < 8:
            target["hp"] -= b["damage"]
            if target["hp"] <= 0:
                gold += target["reward"]
                enemies.remove(target)
            bullets.pop(i)
            continue

        b["x"] += (dx / dist) * b["speed"]
        b["y"] += (dy / dist) * b["speed"]


def check_wave_completion():
    global wave_in_progress, wave, wave_cooldown
    if wave_in_progress and enemies_spawned >= enemies_this_wave and len(enemies) == 0:
        wave_in_progress = False
        wave += 1
        wave_cooldown = 90
        print(f"[波次] 第 {wave - 1} 波完成! 准备下一波...")


def try_build_tower(col, row):
    global gold
    if gold < TOWER_COST:
        return False
    if (col, row) in path_grid:
        return False

    cx = col * GRID_SIZE + GRID_SIZE / 2
    cy = row * GRID_SIZE + GRID_SIZE / 2

    for t in towers:
        if abs(t["x"] - cx) < 10 and abs(t["y"] - cy) < 10:
            return False

    towers.append({
        "x": cx,
        "y": cy,
        "cooldown": 0,
    })
    gold -= TOWER_COST
    return True


def reset_game():
    global lives, gold, wave, towers, enemies, bullets
    global wave_in_progress, enemies_this_wave, enemies_spawned
    global wave_cooldown, game_over
    lives = 10
    gold = 250
    wave = 1
    towers.clear()
    enemies.clear()
    bullets.clear()
    wave_in_progress = False
    enemies_this_wave = 0
    enemies_spawned = 0
    wave_cooldown = 0
    game_over = False


# ==================== 绘制函数 ====================
def draw_grid():
    for i in range(COLS + 1):
        pygame.draw.line(screen, COLOR_GRID, (i * GRID_SIZE, 0),
                         (i * GRID_SIZE, HEIGHT), 1)
    for i in range(ROWS + 1):
        pygame.draw.line(screen, COLOR_GRID, (0, i * GRID_SIZE),
                         (WIDTH, i * GRID_SIZE), 1)


def draw_path():
    pygame.draw.lines(screen, COLOR_PATH, False, waypoints, 34)
    pygame.draw.lines(screen, COLOR_PATH_LIGHT, False, waypoints, 20)


def draw_blocked_cells():
    for (col, row) in path_grid:
        s = pygame.Surface((GRID_SIZE, GRID_SIZE), pygame.SRCALPHA)
        s.fill((180, 70, 70, 40))
        screen.blit(s, (col * GRID_SIZE, row * GRID_SIZE))


def draw_towers():
    for t in towers:
        pygame.draw.circle(screen, COLOR_TOWER_BASE, (int(t["x"]), int(t["y"])), 17)
        pygame.draw.circle(screen, COLOR_TOWER_TOP, (int(t["x"]), int(t["y"])), 12)
        pygame.draw.circle(screen, (224, 242, 255), (int(t["x"]), int(t["y"])), 12, 2)
        pygame.draw.circle(screen, COLOR_TOWER_GUN, (int(t["x"]), int(t["y"])), 6)


def draw_enemies():
    for e in enemies:
        pygame.draw.rect(screen, COLOR_HP_BG,
                         (e["x"] - 16, e["y"] - 22, 32, 5))
        hp_ratio = e["hp"] / e["max_hp"]
        if hp_ratio > 0.6:
            hp_color = COLOR_HP_GREEN
        elif hp_ratio > 0.3:
            hp_color = COLOR_HP_YELLOW
        else:
            hp_color = COLOR_HP_RED
        pygame.draw.rect(screen, hp_color,
                         (e["x"] - 16, e["y"] - 22, 32 * hp_ratio, 5))

        pygame.draw.circle(screen, COLOR_ENEMY,
                           (int(e["x"]), int(e["y"])), e["size"])
        pygame.draw.circle(screen, (255, 255, 255),
                           (int(e["x"] - 4), int(e["y"] - 4)), 3)
        pygame.draw.circle(screen, (255, 255, 255),
                           (int(e["x"] + 4), int(e["y"] - 4)), 3)
        pygame.draw.circle(screen, (15, 15, 15),
                           (int(e["x"] - 4), int(e["y"] - 4)), 1)
        pygame.draw.circle(screen, (15, 15, 15),
                           (int(e["x"] + 4), int(e["y"] - 4)), 1)


def draw_bullets():
    for b in bullets:
        pygame.draw.circle(screen, COLOR_BULLET, (int(b["x"]), int(b["y"])), 5)


def draw_mouse_hint():
    if mouse_col < 0 or mouse_row < 0:
        return
    x = mouse_col * GRID_SIZE
    y = mouse_row * GRID_SIZE
    can_build = (mouse_col, mouse_row) not in path_grid and gold >= TOWER_COST

    cx = mouse_col * GRID_SIZE + GRID_SIZE / 2
    cy = mouse_row * GRID_SIZE + GRID_SIZE / 2
    occupied = any(abs(t["x"] - cx) < 10 and abs(t["y"] - cy) < 10 for t in towers)

    if can_build and not occupied:
        s = pygame.Surface((GRID_SIZE, GRID_SIZE), pygame.SRCALPHA)
        s.fill((100, 255, 150, 90))
        screen.blit(s, (x, y))
        pygame.draw.rect(screen, (170, 255, 170), (x + 1, y + 1, GRID_SIZE - 2, GRID_SIZE - 2), 2)
    else:
        s = pygame.Surface((GRID_SIZE, GRID_SIZE), pygame.SRCALPHA)
        s.fill((255, 80, 80, 60))
        screen.blit(s, (x, y))


def draw_ui():
    pygame.draw.rect(screen, (30, 47, 58), (0, 0, WIDTH, 40))
    pygame.draw.line(screen, (60, 90, 110), (0, 40), (WIDTH, 40), 2)

    lives_text = FONT.render(f"生命: {lives}", True, COLOR_TEXT)
    gold_text = FONT.render(f"金币: {gold}", True, COLOR_GOLD)
    wave_text = FONT.render(f"波次: {wave}", True, COLOR_TEXT)

    screen.blit(lives_text, (20, 10))
    screen.blit(gold_text, (220, 10))
    screen.blit(wave_text, (440, 10))

    hint = FONT_SMALL.render("点击空地建塔 (50金) | 塔自动攻击附近敌人", True, (140, 180, 200))
    screen.blit(hint, (600, 12))


def draw_game_over():
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 180))
    screen.blit(overlay, (0, 0))

    over_text = FONT_BIG.render("游戏结束", True, (255, 100, 100))
    tip_text = FONT.render("按 R 键重新开始", True, (255, 255, 255))

    screen.blit(over_text, (WIDTH // 2 - over_text.get_width() // 2, HEIGHT // 2 - 60))
    screen.blit(tip_text, (WIDTH // 2 - tip_text.get_width() // 2, HEIGHT // 2 + 20))


# ==================== 主循环 ====================
def main():
    global mouse_col, mouse_row, spawn_timer, wave_cooldown

    running = True
    start_wave()

    while running:
        clock.tick(FPS)

        # ---------- 事件处理 ----------
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

            elif event.type == pygame.MOUSEMOTION:
                mx, my = event.pos
                col = mx // GRID_SIZE
                row = my // GRID_SIZE
                if 0 <= col < COLS and 0 <= row < ROWS:
                    mouse_col, mouse_row = col, row
                else:
                    mouse_col, mouse_row = -1, -1

            elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                if not game_over and mouse_col >= 0 and mouse_row >= 0:
                    try_build_tower(mouse_col, mouse_row)

            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r and game_over:
                    reset_game()
                    start_wave()

        # ---------- 更新逻辑 ----------
        if not game_over:
            if not wave_in_progress and wave_cooldown <= 0:
                start_wave()
            if wave_cooldown > 0:
                wave_cooldown -= 1

            if wave_in_progress and enemies_spawned < enemies_this_wave:
                if spawn_timer <= 0:
                    spawn_enemy()
                    spawn_timer = SPAWN_DELAY
                else:
                    spawn_timer -= 1

            update_enemies()
            update_towers()
            update_bullets()
            check_wave_completion()

        # ---------- 绘制 ----------
        screen.fill(COLOR_BG)
        draw_grid()
        draw_path()
        draw_blocked_cells()
        draw_towers()
        draw_enemies()
        draw_bullets()
        draw_mouse_hint()
        draw_ui()

        if game_over:
            draw_game_over()

        pygame.display.flip()

    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()