import pygame
import math
import random
import sys

# ---------- 初始化 ----------
pygame.init()
WIDTH, HEIGHT = 900, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("塔防 - Pygame")
clock = pygame.time.Clock()

# ---- 安全字体初始化 ----
try:
    font = pygame.font.SysFont("simhei", 22, bold=True)
except Exception:
    font = pygame.font.Font(None, 22)

try:
    big_font = pygame.font.SysFont("simhei", 34, bold=True)
except Exception:
    big_font = pygame.font.Font(None, 34)

# ---------- 颜色 ----------
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GRAY = (128, 128, 128)
DARK_GREEN = (0, 100, 0)
LIGHT_BROWN = (210, 170, 110)
RED = (220, 60, 60)
BLUE = (60, 130, 230)
GOLD = (255, 215, 0)
ORANGE = (255, 165, 0)
PURPLE = (138, 43, 226)

# ---------- 游戏状态 ----------
game_state = "menu"
money = 200
lives = 20
wave = 0
max_waves = 7
selected_tower_type = None
placing_tower = False
hover_pos = None

# ---------- 路径点 ----------
path_points = [
    (0, 350), (100, 320), (200, 250), (280, 190), (380, 220),
    (450, 310), (520, 390), (610, 420), (720, 360), (820, 330),
    (900, 340)
]

# ---------- 塔数据 ----------
TOWER_DATA = {
    "arrow":  {"cost": 50,  "range": 140, "damage": 25, "cooldown": 35, "color": (0,180,0)},
    "sniper": {"cost": 120, "range": 260, "damage": 90, "cooldown": 75, "color": (200,50,50)},
    "ice":    {"cost": 80,  "range": 150, "damage": 15, "cooldown": 45, "color": (100,200,255)},
    "cannon": {"cost": 100, "range": 135, "damage": 55, "cooldown": 65, "color": (200,100,0)}
}

# ---------- 精灵组 ----------
all_sprites = pygame.sprite.Group()
enemies = pygame.sprite.Group()
towers = pygame.sprite.Group()
projectiles = pygame.sprite.Group()

# ---------- 辅助函数 ----------
def point_to_line_distance(point, start, end):
    px, py = point
    sx, sy = start
    ex, ey = end
    dx = ex - sx
    dy = ey - sy
    if dx == 0 and dy == 0:
        return math.hypot(px - sx, py - sy)
    t = ((px - sx)*dx + (py - sy)*dy) / (dx*dx + dy*dy)
    t = max(0, min(1, t))
    nx = sx + t * dx
    ny = sy + t * dy
    return math.hypot(px - nx, py - ny)

def can_build(x, y):
    for i in range(len(path_points)-1):
        if point_to_line_distance((x, y), path_points[i], path_points[i+1]) < 26:
            return False
    for t in towers:
        if math.hypot(x - t.rect.centerx, y - t.rect.centery) < 44:
            return False
    return True

# ---------- 敌人 ----------
class Enemy(pygame.sprite.Sprite):
    def __init__(self, path_idx=0, hp=50, speed=2, reward=20):
        super().__init__()
        self.hp_max = hp
        self.hp = hp
        self.speed = speed
        self.reward = reward
        self.path_index = path_idx
        self.next_point = path_points[self.path_index]
        self.prev_point = path_points[0]
        self.pos = [float(self.prev_point[0]), float(self.prev_point[1])]
        self.image = pygame.Surface((18, 14))
        self.image.fill(RED)
        self.rect = self.image.get_rect(center=(int(self.pos[0]), int(self.pos[1])))
        self.frozen_timer = 0

    def update(self):
        global lives, money
        if self.frozen_timer > 0:
            self.frozen_timer -= 1
            speed_mult = 0.4
        else:
            speed_mult = 1.0

        tx, ty = self.next_point
        dx = tx - self.pos[0]
        dy = ty - self.pos[1]
        dist = math.hypot(dx, dy)
        if dist < 2:
            self.path_index += 1
            if self.path_index >= len(path_points):
                lives -= 1
                self.kill()
                return
            self.prev_point = self.next_point
            self.next_point = path_points[self.path_index]
        else:
            step = self.speed * speed_mult
            self.pos[0] += (dx / dist) * step
            self.pos[1] += (dy / dist) * step
            self.rect.center = (int(self.pos[0]), int(self.pos[1]))

        if self.hp <= 0:
            money += self.reward
            self.kill()

    def take_damage(self, damage, slow=False):
        self.hp -= damage
        if slow:
            self.frozen_timer = 60

# ---------- 防御塔 ----------
class Tower(pygame.sprite.Sprite):
    def __init__(self, x, y, tower_type):
        super().__init__()
        self.tower_type = tower_type
        data = TOWER_DATA[tower_type]
        self.level = 1
        self.damage = data["damage"]
        self.range = data["range"]
        self.cooldown = data["cooldown"]
        self.color = data["color"]
        self.cost = data["cost"]
        self.image = pygame.Surface((32, 32))
        self.image.fill(self.color)
        self.rect = self.image.get_rect(center=(x, y))
        self.cooldown_timer = 0
        self.target = None

    def upgrade(self):
        global money
        if self.level < 5 and money >= self.cost * 0.6:
            money -= int(self.cost * 0.6)
            self.level += 1
            self.damage = int(self.damage * 1.5)
            self.range = int(self.range * 1.1)
            self.cooldown = max(10, self.cooldown - 3)
            return True
        return False

    def sell_value(self):
        total_invested = self.cost + sum(int(self.cost * 0.6) for _ in range(self.level-1))
        return int(total_invested * 0.5)

    def find_target(self):
        best = None
        best_dist = float('inf')
        for e in enemies:
            d = math.hypot(self.rect.centerx - e.rect.centerx,
                           self.rect.centery - e.rect.centery)
            if d < self.range and d < best_dist:
                best = e
                best_dist = d
        return best

    def update(self):
        self.cooldown_timer = max(0, self.cooldown_timer - 1)
        if self.cooldown_timer == 0:
            target = self.find_target()
            if target:
                self.cooldown_timer = self.cooldown
                bullet = Bullet(self.rect.centerx, self.rect.centery,
                                target, self.damage, self.tower_type)
                all_sprites.add(bullet)
                projectiles.add(bullet)

# ---------- 子弹 ----------
class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y, target, damage, tower_type):
        super().__init__()
        self.image = pygame.Surface((6, 6))
        if tower_type == "ice":
            self.image.fill((200, 235, 255))
        elif tower_type == "cannon":
            self.image.fill((255, 140, 0))
        else:
            self.image.fill((255, 255, 0))
        self.rect = self.image.get_rect(center=(x, y))
        self.target = target
        self.damage = damage
        self.tower_type = tower_type
        self.speed = 8

    def update(self):
        if not self.target.alive():
            self.kill()
            return
        dx = self.target.rect.centerx - self.rect.centerx
        dy = self.target.rect.centery - self.rect.centery
        dist = math.hypot(dx, dy)
        if dist < 10:
            slow = (self.tower_type == "ice")
            self.target.take_damage(self.damage, slow)
            if self.tower_type == "cannon":
                splash_range = 42
                for e in enemies:
                    if e != self.target and e.alive():
                        d = math.hypot(e.rect.centerx - self.target.rect.centerx,
                                       e.rect.centery - self.target.rect.centery)
                        if d < splash_range:
                            e.take_damage(int(self.damage * 0.4))
            self.kill()
        else:
            step = self.speed
            self.rect.x += int((dx / dist) * step)
            self.rect.y += int((dy / dist) * step)

# ---------- 生成敌人 ----------
def spawn_wave(wave_number):
    wave_config = {
        1: {"count": 5,  "hp": 40,  "speed": 2.0, "reward": 15},
        2: {"count": 8,  "hp": 50,  "speed": 2.2, "reward": 16},
        3: {"count": 10, "hp": 65,  "speed": 2.3, "reward": 17},
        4: {"count": 12, "hp": 85,  "speed": 2.5, "reward": 18},
        5: {"count": 15, "hp": 105, "speed": 2.7, "reward": 20},
        6: {"count": 18, "hp": 125, "speed": 2.9, "reward": 23},
        7: {"count": 21, "hp": 155, "speed": 3.1, "reward": 27},
    }
    cfg = wave_config.get(wave_number, wave_config[7])
    for _ in range(cfg["count"]):
        enemy = Enemy(path_idx=1, hp=cfg["hp"], speed=cfg["speed"], reward=cfg["reward"])
        all_sprites.add(enemy)
        enemies.add(enemy)

# ---------- UI ----------
def draw_ui():
    info_y = 650
    pygame.draw.rect(screen, (30,30,30), (0, info_y, WIDTH, 50))
    texts = [
        f"💰 ${money}",
        f"❤️ {lives}",
        f"🌊 波次 {wave}/{max_waves}"
    ]
    for i, txt in enumerate(texts):
        surf = font.render(txt, True, WHITE)
        screen.blit(surf, (20 + i*200, info_y+12))

    panel_x = 745
    panel_y = 10
    pygame.draw.rect(screen, (40,40,40), (panel_x-10, panel_y-10, 156, 272))
    types = ["arrow", "sniper", "ice", "cannon"]
    names = ["箭塔", "狙击", "冰塔", "炮塔"]
    costs = [50, 120, 80, 100]
    for i, (t, n, c) in enumerate(zip(types, names, costs)):
        y = panel_y + i*63
        color = TOWER_DATA[t]["color"]
        pygame.draw.rect(screen, color, (panel_x, y, 33, 33))
        label = font.render(f"{n} ${c}", True, WHITE)
        screen.blit(label, (panel_x+39, y+6))
        if selected_tower_type == t:
            pygame.draw.rect(screen, GOLD, (panel_x-3, y-3, 146, 41), 2)

    hint = font.render("1-4选塔, 左键放置, ESC取消", True, GRAY)
    screen.blit(hint, (10, 612))

def draw_menu():
    screen.fill((20, 40, 71))
    title = big_font.render("塔 防 游 戏", True, GOLD)
    screen.blit(title, (WIDTH//2 - 121, 191))
    start = font.render("按 空格键 开始游戏", True, WHITE)
    screen.blit(start, (WIDTH//2 - 126, 281))
    instr = font.render("方向键选塔, 左键放置, 右键取消", True, GRAY)
    screen.blit(instr, (WIDTH//2 - 181, 331))

def draw_game_over():
    overlay = pygame.Surface((WIDTH, HEIGHT))
    overlay.set_alpha(192)
    overlay.fill((0,0,0))
    screen.blit(overlay, (0,0))
    text = big_font.render("游 戏 结 束", True, RED)
    screen.blit(text, (WIDTH//2 - 123, 271))
    score = font.render(f"坚持到第 {wave} 波 | 剩余金币 ${money}", True, WHITE)
    screen.blit(score, (WIDTH//2 - 163, 341))
    restart = font.render("按 R 键重新开始", True, GOLD)
    screen.blit(restart, (WIDTH//2 - 111, 401))

# ---------- 重置 ----------
def reset_game():
    global money, lives, wave, selected_tower_type, placing_tower, game_state
    global all_sprites, enemies, towers, projectiles
    money = 200
    lives = 20
    wave = 0
    selected_tower_type = None
    placing_tower = False
    game_state = "playing"
    all_sprites.empty()
    enemies.empty()
    towers.empty()
    projectiles.empty()

# ---------- 主循环 ----------
running = True
while running:
    dt = clock.tick(60)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        if game_state == "menu":
            if event.type == pygame.KEYDOWN and event.key == pygame.K_SPACE:
                reset_game()
                game_state = "playing"

        elif game_state == "playing":
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_1:
                    selected_tower_type = "arrow"
                elif event.key == pygame.K_2:
                    selected_tower_type = "sniper"
                elif event.key == pygame.K_3:
                    selected_tower_type = "ice"
                elif event.key == pygame.K_4:
                    selected_tower_type = "cannon"
                elif event.key == pygame.K_ESCAPE:
                    selected_tower_type = None
                    placing_tower = False
                elif event.key == pygame.K_SPACE:
                    if wave < max_waves and len(enemies) == 0:
                        wave += 1
                        spawn_wave(wave)

            if event.type == pygame.MOUSEMOTION:
                hover_pos = event.pos

            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                mx, my = event.pos
                if 735 <= mx <= 891 and 0 <= my <= 254:
                    idx = (my - 10) // 68
                    if 0 <= idx < 4:
                        selected_tower_type = ["arrow","sniper","ice","cannon"][idx]
                else:
                    if selected_tower_type and my < 651:
                        cost = TOWER_DATA[selected_tower_type]["cost"]
                        if money >= cost and can_build(mx, my):
                            money -= cost
                            new_tower = Tower(mx, my, selected_tower_type)
                            all_sprites.add(new_tower)
                            towers.add(new_tower)
                        else:
                            for t in towers:
                                if t.rect.collidepoint(mx, my):
                                    if event.button == 1:
                                        t.upgrade()
                                    break

            if event.type == pygame.MOUSEBUTTONDOWN and event.button == 3:
                mx, my = event.pos
                for t in towers:
                    if t.rect.collidepoint(mx, my):
                        money += t.sell_value()
                        t.kill()
                        break
                else:
                    selected_tower_type = None
                    placing_tower = False

        elif game_state == "gameover":
            if event.type == pygame.KEYDOWN and event.key == pygame.K_r:
                reset_game()
                game_state = "playing"

    if game_state == "playing":
        all_sprites.update()
        if lives <= 0:
            game_state = "gameover"

    screen.fill(DARK_GREEN)
    for i in range(len(path_points)-1):
        pygame.draw.line(screen, LIGHT_BROWN, path_points[i], path_points[i+1], 18)
    for pt in path_points:
        pygame.draw.circle(screen, (139,90,43), pt, 4)

    if selected_tower_type and hover_pos:
        mx, my = hover_pos
        if can_build(mx, my):
            r = TOWER_DATA[selected_tower_type]["range"]
            pygame.draw.circle(screen, (255,255,255,40), (mx, my), r, 1)

    all_sprites.draw(screen)

    if game_state == "menu":
        draw_menu()
    elif game_state == "playing":
        draw_ui()
    elif game_state == "gameover":
        draw_game_over()

    pygame.display.flip()

pygame.quit()
sys.exit()