import pygame
import sys
import math
import random

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

WIDTH, HEIGHT = 1000, 700
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Radish Defense - Tower Defense")

# ===== 颜色（所有值严格 0~255） =====
BG          = (35, 55, 35)
PATH_COLOR  = (175, 155, 95)
GRASS       = (65, 135, 65)
TOWER_COLORS = {
    "basic":   (170, 170, 190),
    "sniper":  (90, 190, 90),
    "machine": (190, 95, 95),
    "cannon":  (195, 145, 45),
    "freeze":  (90, 170, 250),
}
BULLET      = (255, 200, 0)
ENEMY       = (215, 75, 75)
HP_GREEN    = (75, 215, 75)
HP_RED      = (215, 75, 75)
UI_BG       = (40, 40, 60)
UI_TEXT     = (215, 215, 250)
BTN_COLOR   = (65, 95, 145)
BTN_HOVER   = (85, 125, 185)
RADISH_BODY = (235, 115, 155)
RADISH_LVS  = (95, 215, 95)
SLOW_CLR    = (95, 195, 250)
EXPLODE     = (255, 145, 0)
WHITE       = (255, 255, 255)
BLACK       = (0, 0, 0)
YELLOW      = (255, 210, 30)
GRID_OK     = (60, 120, 60, 100)       # 可放置区域（半透明绿）
GRID_OK_BD  = (80, 160, 80, 140)
GRID_NO     = (160, 60, 60, 90)        # 不可放置（半透明红）
GRID_NO_BD  = (200, 80, 80, 120)

FPS = 60
clock = pygame.time.Clock()
font = pygame.font.Font(None, 24)
big_font = pygame.font.Font(None, 48)
title_font = pygame.font.Font(None, 72)

# ===== 路径 =====
path = [
    (50, 350), (240, 350), (240, 210), (420, 210),
    (420, 410), (600, 410), (630, 300), (820, 310), (870, 510)
]

# ===== 可放置区域（用网格标记） =====
CELL = 50  # 网格大小
cols, rows = WIDTH // CELL, (HEIGHT - 150) // CELL  # 底部留150给UI

def point_to_cell(px, py):
    return px // CELL, py // CELL

def cell_center(cx, cy):
    return cx * CELL + CELL // 2, cy * CELL + CELL // 2

def is_on_path(px, py, margin=28):
    """判断点是否在路径附近"""
    for i in range(len(path) - 1):
        x1, y1 = path[i]
        x2, y2 = path[i+1]
        seg_len = math.sqrt((x2-x1)**2 + (y2-y1)**2)
        if seg_len == 0:
            continue
        # 点到线段的距离
        t = max(0, min(1, ((px-x1)*(x2-x1) + (py-y1)*(y2-y1)) / (seg_len**2)))
        proj_x = x1 + t * (x2 - x1)
        proj_y = y1 + t * (y2 - y1)
        d = math.sqrt((px - proj_x)**2 + (py - proj_y)**2)
        if d < margin:
            return True
    return False

def is_near_radish(px, py, margin=45):
    rx, ry = path[-1]
    return math.sqrt((px - rx)**2 + (py - ry)**2) < margin

def is_near_tower(px, py, towers, margin=45):
    for t in towers:
        if math.sqrt((px - t.x)**2 + (py - t.y)**2) < margin:
            return True
    return False

def build_grid(towers):
    """构建可放置网格标记"""
    grid = {}
    for cy in range(rows):
        for cx in range(cols):
            px, py = cell_center(cx, cy)
            if py >= HEIGHT - 150:  # UI区域不可放
                grid[(cx, cy)] = "no"
                continue
            if is_on_path(px, py):
                grid[(cx, cy)] = "no"
            elif is_near_radish(px, py):
                grid[(cx, cy)] = "no"
            elif is_near_tower(px, py, towers):
                grid[(cx, cy)] = "no"
            else:
                grid[(cx, cy)] = "ok"
    return grid

def draw_grid(surface, grid, mouse_pos):
    """绘制网格标记"""
    mx, my = mouse_pos
    for (cx, cy), status in grid.items():
        x = cx * CELL
        y = cy * CELL
        if status == "ok":
            # 可放置：淡绿方块
            s = pygame.Surface((CELL-2, CELL-2), pygame.SRCALPHA)
            s.fill(GRID_OK)
            surface.blit(s, (x+1, y+1))
            pygame.draw.rect(surface, GRID_OK_BD, (x+1, y+1, CELL-2, CELL-2), 1)
            # 中心小绿点
            pygame.draw.circle(surface, (80, 180, 80), (x+CELL//2, y+CELL//2), 3)
        else:
            # 不可放置：淡红斜线
            s = pygame.Surface((CELL-2, CELL-2), pygame.SRCALPHA)
            s.fill(GRID_NO)
            surface.blit(s, (x+1, y+1))
            pygame.draw.rect(surface, GRID_NO_BD, (x+1, y+1, CELL-2, CELL-2), 1)

        # 鼠标悬停的格子高亮
        if (mx // CELL, my // CELL) == (cx, cy):
            if status == "ok":
                pygame.draw.rect(surface, YELLOW, (x, y, CELL, CELL), 3)
            else:
                pygame.draw.rect(surface, (255, 80, 80), (x, y, CELL, CELL), 3)

# ===== 游戏状态 =====
money = 150
health = 10
game_over = False
game_won = False
wave = 1
enemies_defeated = 0
enemies_per_wave = 15
spawn_timer = 0
spawn_delay = 40
selected_tower_type = None
tower_info_visible = False
info_tower = None

# ===== 炮塔类型 =====
TOWER_TYPES = {
    "basic":   {"name": "Basic",   "cost": 60,  "range": 150, "damage": 1,   "cooldown": 30, "color": TOWER_COLORS["basic"]},
    "sniper":  {"name": "Sniper",  "cost": 120, "range": 300, "damage": 3,   "cooldown": 60, "color": TOWER_COLORS["sniper"]},
    "machine": {"name": "Machine", "cost": 100, "range": 120, "damage": 0.5, "cooldown": 10, "color": TOWER_COLORS["machine"]},
    "cannon":  {"name": "Cannon",  "cost": 150, "range": 130, "damage": 2,   "cooldown": 40, "color": TOWER_COLORS["cannon"]},
    "freeze":  {"name": "Freeze",  "cost": 130, "range": 140, "damage": 0.2, "cooldown": 50, "color": TOWER_COLORS["freeze"]},
}


# ===== Tower =====
class Tower:
    def __init__(self, x, y, tower_type):
        self.x = x
        self.y = y
        self.type = tower_type
        self.stats = TOWER_TYPES[tower_type]
        self.range = self.stats["range"]
        self.damage = self.stats["damage"]
        self.cooldown = 0
        self.cooldown_max = self.stats["cooldown"]
        self.cost = self.stats["cost"]
        self.level = 1
        self.color = self.stats["color"]
        self.gun_clr = (max(0, self.color[0]-50), max(0, self.color[1]-50), max(0, self.color[2]-50))

    def draw(self, s):
        pygame.draw.circle(s, self.color, (self.x, self.y), 20)
        darker = (max(0, self.color[0]-30), max(0, self.color[1]-30), max(0, self.color[2]-30))
        pygame.draw.circle(s, darker, (self.x, self.y), 15)
        if self.type == "basic":
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y-5, 25, 10))
        elif self.type == "sniper":
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y-3, 35, 6))
            pygame.draw.circle(s, (50, 50, 50), (self.x+40, self.y), 5)
        elif self.type == "machine":
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y-8, 25, 6))
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y+2, 25, 6))
        elif self.type == "cannon":
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y-7, 30, 14))
        elif self.type == "freeze":
            pygame.draw.rect(s, self.gun_clr, (self.x+5, self.y-5, 25, 10))
            pygame.draw.circle(s, (195, 225, 255), (self.x+15, self.y), 8, 2)

    def update(self, enemies):
        if self.cooldown > 0:
            self.cooldown -= 1
            return None
        for e in enemies:
            if not e.active or e.health <= 0:
                continue
            d = math.sqrt((self.x - e.x)**2 + (self.y - e.y)**2)
            if d <= self.range:
                self.cooldown = self.cooldown_max
                return Bullet(self.x, self.y, e, self.damage, self.type)
        return None


# ===== Bullet（修复：所有分支都有 self.size） =====
class Bullet:
    def __init__(self, x, y, target, damage, btype):
        self.x = x
        self.y = y
        self.target = target
        self.speed = 8
        self.damage = damage
        self.active = True
        self.type = btype
        self.effect = None
        self.size = 5  # ★ 默认大小，所有子弹都有

        if btype == "freeze":
            self.color = (100, 200, 255)
            self.effect = "slow"
            self.effect_duration = 180
        elif btype == "cannon":
            self.color = (255, 150, 0)
            self.size = 8  # 炮弹更大
        elif btype == "machine":
            self.color = BULLET
            self.size = 3  # 机枪子弹更小更快
        elif btype == "sniper":
            self.color = (150, 255, 150)
            self.size = 3
        else:  # basic
            self.color = BULLET
            # size 保持默认 5

    def update(self):
        if not self.active or not self.target.active:
            self.active = False
            return False
        dx = self.target.x - self.x
        dy = self.target.y - self.y
        dist = math.sqrt(dx**2 + dy**2)
        if dist < self.speed:
            self.target.health -= self.damage
            if self.effect == "slow":
                self.target.slow_duration = self.effect_duration
                self.target.speed = self.target.original_speed * 0.5
            self.active = False
            return True
        self.x += dx / dist * self.speed
        self.y += dy / dist * self.speed
        return False

    def draw(self, s):
        if self.type == "cannon":
            pygame.draw.circle(s, self.color, (int(self.x), int(self.y)), self.size)
            pygame.draw.circle(s, (255, 200, 100), (int(self.x), int(self.y)), max(1, self.size-3))
        else:
            pygame.draw.circle(s, self.color, (int(self.x), int(self.y)), self.size)
            pygame.draw.circle(s, (255, 255, 200), (int(self.x), int(self.y)), max(1, self.size-2))


# ===== Enemy =====
class Enemy:
    def __init__(self, path, wave):
        self.path = path
        self.path_index = 0
        self.x, self.y = path[0]
        self.max_health = 5 + (wave - 1) * 2
        self.health = self.max_health
        self.speed = 1.0 + (wave - 1) * 0.1
        self.original_speed = self.speed
        self.reward = 10 + (wave - 1) * 2
        self.slow_duration = 0
        self.active = True
        r = random.random()
        if r < 0.25:
            self.type = "fast"
            self.speed *= 1.8
            self.health *= 0.7
            self.reward = int(self.reward * 1.3)
            self.color = (220, 120, 120)
        elif r < 0.45:
            self.type = "tank"
            self.speed *= 0.7
            self.health *= 3.0
            self.reward = int(self.reward * 1.8)
            self.color = (180, 80, 80)
        else:
            self.type = "normal"
            self.color = ENEMY

    def update(self):
        if not self.active:
            return False
        if self.slow_duration > 0:
            self.slow_duration -= 1
            if self.slow_duration == 0:
                self.speed = self.original_speed
        tx, ty = self.path[self.path_index]
        dx = tx - self.x
        dy = ty - self.y
        dist = math.sqrt(dx**2 + dy**2)
        if dist < self.speed:
            self.path_index += 1
            if self.path_index >= len(self.path):
                self.active = False
                return True
            tx, ty = self.path[self.path_index]
            dx = tx - self.x
            dy = ty - self.y
            dist = math.sqrt(dx**2 + dy**2)
        self.x += dx / dist * self.speed
        self.y += dy / dist * self.speed
        return False

    def draw(self, s):
        pygame.draw.circle(s, self.color, (int(self.x), int(self.y)), 15)
        darker = (max(0, self.color[0]-40), max(0, self.color[1]-40), max(0, self.color[2]-40))
        if self.type == "fast":
            pygame.draw.circle(s, darker, (int(self.x), int(self.y)), 10)
            pygame.draw.line(s, (255, 255, 200), (self.x, self.y-15), (self.x, self.y-25), 2)
        elif self.type == "tank":
            d2 = (max(0, self.color[0]-50), max(0, self.color[1]-50), max(0, self.color[2]-50))
            pygame.draw.rect(s, d2, (self.x-12, self.y-5, 24, 10))
            pygame.draw.rect(s, d2, (self.x-5, self.y-12, 10, 24))
        else:
            pygame.draw.circle(s, darker, (int(self.x), int(self.y)), 10)
        bw = 40
        pygame.draw.rect(s, HP_RED, (self.x - bw//2, self.y - 30, bw, 5))
        pygame.draw.rect(s, HP_GREEN, (self.x - bw//2, self.y - 30, int(bw * self.health / self.max_health), 5))
        if self.slow_duration > 0:
            pygame.draw.circle(s, SLOW_CLR, (int(self.x), int(self.y)), 18, 2)


# ===== Explosion =====
class Explosion:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = 5
        self.max_radius = 30
        self.growth = 1.5
        self.active = True

    def update(self):
        self.radius += self.growth
        if self.radius > self.max_radius:
            self.active = False

    def draw(self, s):
        a = int(255 * (1 - self.radius / self.max_radius))
        surf = pygame.Surface((self.max_radius*2, self.max_radius*2), pygame.SRCALPHA)
        pygame.draw.circle(surf, (*EXPLODE, a), (self.max_radius, self.max_radius), int(self.radius), 3)
        pygame.draw.circle(surf, (255, 255, 200, a//2), (self.max_radius, self.max_radius), int(self.radius/2), 2)
        s.blit(surf, (int(self.x - self.max_radius), int(self.y - self.max_radius)))


# ===== 游戏对象列表 =====
towers = []
bullets = []
enemies = []
explosions = []


# ===== 绘图函数 =====
def draw_path_line():
    for i in range(len(path) - 1):
        pygame.draw.line(screen, PATH_COLOR, path[i], path[i+1], 40)
    # 起点装饰
    pygame.draw.circle(screen, (115, 85, 55), path[0], 20)
    pygame.draw.circle(screen, (140, 110, 75), path[0], 14)


def draw_grass_bg():
    for i in range(0, WIDTH, 40):
        for j in range(0, HEIGHT - 150, 40):
            if random.random() > 0.35:
                pygame.draw.line(screen, GRASS, (i, j), (i, j-14), 2)


def draw_radish():
    x, y = path[-1]
    # 叶子
    pygame.draw.ellipse(screen, RADISH_LVS, (x-24, y-68, 48, 38))
    ld = (max(0, RADISH_LVS[0]-20), max(0, RADISH_LVS[1]-20), max(0, RADISH_LVS[2]-20))
    pygame.draw.ellipse(screen, ld, (x-30, y-58, 28, 28))
    pygame.draw.ellipse(screen, ld, (x+2, y-58, 28, 28))
    # 身体
    pygame.draw.circle(screen, RADISH_BODY, (x, y), 30)
    rd = (max(0, RADISH_BODY[0]-40), max(0, RADISH_BODY[1]-40), max(0, RADISH_BODY[2]-40))
    pygame.draw.ellipse(screen, rd, (x-19, y-14, 38, 28))
    # 眼睛
    pygame.draw.circle(screen, WHITE, (x-10, y-5), 8)
    pygame.draw.circle(screen, WHITE, (x+10, y-5), 8)
    pygame.draw.circle(screen, BLACK, (x-10, y-5), 4)
    pygame.draw.circle(screen, BLACK, (x+10, y-5), 4)
    # 嘴巴
    pygame.draw.arc(screen, (195, 75, 115), (x-10, y+2, 20, 14), 0, math.pi, 3)
    # 血条
    pygame.draw.rect(screen, HP_RED, (x-30, y-95, 60, 8))
    pygame.draw.rect(screen, HP_GREEN, (x-30, y-95, int(60 * health / 10), 8))


def draw_ui():
    ui_y = HEIGHT - 150
    pygame.draw.rect(screen, UI_BG, (0, 0, WIDTH, 48))
    pygame.draw.rect(screen, (50, 50, 70), (0, ui_y, WIDTH, 150))
    pygame.draw.line(screen, (80, 80, 100), (0, ui_y), (WIDTH, ui_y), 3)

    # 顶部信息
    items = [
        f"Gold: {money}",
        f"HP: {health}",
        f"Wave: {wave}/10",
        f"Killed: {enemies_defeated}/{wave * enemies_per_wave}"
    ]
    xs = [20, 160, 310, 470]
    for txt, xp in zip(items, xs):
        t = font.render(txt, True, UI_TEXT)
        screen.blit(t, (xp, 14))

    # 炮塔按钮
    label = font.render("Towers:", True, UI_TEXT)
    screen.blit(label, (20, ui_y + 15))

    bw, bh, sp = 145, 100, 18
    sx = 20
    for i, (tid, td) in enumerate(TOWER_TYPES.items()):
        x = sx + i * (bw + sp)
        y = ui_y + 42
        r = pygame.Rect(x, y, bw, bh)
        if selected_tower_type == tid:
            c = (max(0, td["color"][0]//2), max(0, td["color"][1]//2), max(0, td["color"][2]//2))
        elif r.collidepoint(pygame.mouse.get_pos()):
            c = (max(0, int(td["color"][0]/1.2)), max(0, int(td["color"][1]/1.2)), max(0, int(td["color"][2]/1.2)))
        else:
            c = td["color"]
        pygame.draw.rect(screen, c, r)
        bc = (max(0, c[0]//2), max(0, c[1]//2), max(0, c[2]//2))
        pygame.draw.rect(screen, bc, r, 3)
        # 名字
        nt = font.render(td["name"], True, WHITE)
        screen.blit(nt, (x + bw//2 - nt.get_width()//2, y + 8))
        # 价格
        ct = font.render(f"${td['cost']}", True, (255, 255, 180))
        screen.blit(ct, (x + bw//2 - ct.get_width()//2, y + 38))
        # 图标
        ic = (max(0, c[0]-50), max(0, c[1]-50), max(0, c[2]-50))
        pygame.draw.circle(screen, ic, (x + bw//2, y + 75), 14)
        if tid == "sniper":
            pygame.draw.rect(screen, (50, 50, 50), (x + bw//2 + 4, y + 75 - 3, 18, 5))
        elif tid == "machine":
            pygame.draw.rect(screen, (50, 50, 50), (x + bw//2 + 4, y + 75 - 7, 18, 5))
            pygame.draw.rect(screen, (50, 50, 50), (x + bw//2 + 4, y + 75 + 2, 18, 5))
        elif tid == "cannon":
            pygame.draw.rect(screen, (50, 50, 50), (x + bw//2 + 4, y + 75 - 6, 22, 12))
        elif tid == "freeze":
            pygame.draw.circle(screen, (195, 225, 255), (x + bw//2, y + 75), 9, 2)


def draw_tower_info_panel():
    if not tower_info_visible or not info_tower:
        return
    pw, ph = 300, 200
    px, py = WIDTH//2 - pw//2, HEIGHT//2 - ph//2
    surf = pygame.Surface((pw, ph), pygame.SRCALPHA)
    surf.fill((30, 30, 55, 225))
    pygame.draw.rect(surf, (100, 100, 150), (0, 0, pw, ph), 3)
    t = font.render(f"{TOWER_TYPES[info_tower.type]['name']} Info", True, (255, 255, 200))
    surf.blit(t, (pw//2 - t.get_width()//2, 15))
    for i, line in enumerate([
        f"Damage: {info_tower.damage}",
        f"Range:  {info_tower.range}",
        f"CD:     {info_tower.cooldown_max}F",
        f"DPS:    {info_tower.damage / (info_tower.cooldown_max/60):.1f}" if info_tower.cooldown_max > 0 else "DPS: N/A",
    ]):
        lt = font.render(line, True, UI_TEXT)
        surf.blit(lt, (30, 50 + i * 28))
    if info_tower.type == "freeze":
        ft = font.render("Effect: Slow 50%", True, SLOW_CLR)
        surf.blit(ft, (30, 158))
    # 关闭按钮
    pygame.draw.rect(surf, (180, 75, 75), (pw - 50, 10, 40, 22))
    xt = font.render("X", True, WHITE)
    surf.blit(xt, (pw - 37, 13))
    screen.blit(surf, (px, py))


def draw_game_over_screen():
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 185))
    screen.blit(overlay, (0, 0))
    if game_won:
        txt = big_font.render("Victory! Radish is safe!", True, (95, 250, 95))
    else:
        txt = big_font.render("Game Over! Radish eaten!", True, (250, 95, 95))
    screen.blit(txt, (WIDTH//2 - txt.get_width()//2, HEIGHT//2 - 60))
    rt = font.render("Press R to Restart", True, (195, 200, 250))
    screen.blit(rt, (WIDTH//2 - rt.get_width()//2, HEIGHT//2 + 10))
    st = font.render(f"Score: {enemies_defeated * 10 + money}", True, (250, 250, 145))
    screen.blit(st, (WIDTH//2 - st.get_width()//2, HEIGHT//2 + 50))


def draw_start_screen():
    screen.fill((20, 38, 20))
    # 标题
    tt = title_font.render("Radish Defense", True, (95, 215, 95))
    screen.blit(tt, (WIDTH//2 - tt.get_width()//2, 140))
    # 萝卜
    pygame.draw.circle(screen, RADISH_BODY, (WIDTH//2, 285), 58)
    rd = (max(0, RADISH_BODY[0]-40), max(0, RADISH_BODY[1]-40), max(0, RADISH_BODY[2]-40))
    pygame.draw.ellipse(screen, rd, (WIDTH//2-38, 365, 76, 48))
    pygame.draw.circle(screen, WHITE, (WIDTH//2-19, 385), 14)
    pygame.draw.circle(screen, WHITE, (WIDTH//2+19, 393), 14)
    pygame.draw.circle(screen, BLACK, (WIDTH//2-19, 387), 6)
    pygame.draw.circle(screen, BLACK, (WIDTH//2+19, 397), 6)
    pygame.draw.arc(screen, (195, 75, 115), (WIDTH//2-18, 375, 36, 26), 0, math.pi, 5)
    # 叶子
    pygame.draw.ellipse(screen, RADISH_LVS, (WIDTH//2-38, 275, 76, 56))
    ld = (max(0, RADISH_LVS[0]-20), max(0, RADISH_LVS[1]-20), max(0, RADISH_LVS[2]-20))
    pygame.draw.ellipse(screen, ld, (WIDTH//2-58, 302, 36, 36))
    pygame.draw.ellipse(screen, ld, (WIDTH//2+18, 332, 36, 36))
    # 开始按钮
    bx, by, bw, bh = WIDTH//2 - 100, 470, 200, 60
    br = pygame.Rect(bx, by, bw, bh)
    pygame.draw.rect(screen, BTN_COLOR, br)
    if br.collidepoint(pygame.mouse.get_pos()):
        pygame.draw.rect(screen, BTN_HOVER, br, 3)
    else:
        pygame.draw.rect(screen, BTN_HOVER, br, 2)
    stxt = big_font.render("Start", True, WHITE)
    screen.blit(stxt, (WIDTH//2 - stxt.get_width()//2, by + 13))
    # 说明
    tips = [
        "Defend the radish from enemies!",
        "Click a tower type, then click a green grid to place it.",
        "Red grids = blocked.  Green = buildable.",
        "Survive 10 waves to win!",
    ]
    for i, tip in enumerate(tips):
        t = font.render(tip, True, (195, 215, 195))
        screen.blit(t, (WIDTH//2 - t.get_width()//2, 555 + i * 28))


def reset_game():
    global money, health, game_over, game_won, wave, enemies_defeated
    global towers, bullets, enemies, explosions, selected_tower_type
    global spawn_timer, tower_info_visible, info_tower
    money = 150
    health = 10
    game_over = False
    game_won = False
    wave = 1
    enemies_defeated = 0
    towers = []
    bullets = []
    enemies = []
    explosions = []
    selected_tower_type = None
    spawn_timer = 0
    tower_info_visible = False
    info_tower = None


# ===== 主循环 =====
running = True
start_screen = True

while running:
    mouse_pos = pygame.mouse.get_pos()

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

        elif ev.type == pygame.MOUSEBUTTONDOWN:
            if start_screen:
                bx, by, bw, bh = WIDTH//2-100, 470, 200, 60
                if bx <= ev.pos[0] <= bx+bw and by <= ev.pos[1] <= by+bh:
                    start_screen = False
                continue

            if game_over:
                continue

            mx, my = ev.pos

            # 关闭信息面板
            if tower_info_visible:
                px, py = WIDTH//2-150, HEIGHT//2-100
                if px+250 <= mx <= px+290 and py+10 <= my <= py+32:
                    tower_info_visible = False
                continue

            ui_y = HEIGHT - 150

            # 点击 UI 炮塔按钮
            if my >= ui_y:
                bw2, sp = 145, 18
                sx = 20
                for i, tid in enumerate(TOWER_TYPES.keys()):
                    bx2 = sx + i * (bw2 + sp)
                    by2 = ui_y + 42
                    if bx2 <= mx <= bx2+bw2 and by2 <= my <= by2+100:
                        if money >= TOWER_TYPES[tid]["cost"]:
                            selected_tower_type = tid
                        break
                continue

            # 点击地图放置炮塔
            if selected_tower_type and my < ui_y:
                grid = build_grid(towers)
                cell = point_to_cell(mx, my)
                if grid.get(cell) == "ok":
                    cx, cy = cell_center(*cell)
                    # 再检查一次（防边缘情况）
                    if not is_on_path(cx, cy) and not is_near_radish(cx, cy) and not is_near_tower(cx, cy, towers):
                        if money >= TOWER_TYPES[selected_tower_type]["cost"]:
                            towers.append(Tower(cx, cy, selected_tower_type))
                            money -= TOWER_TYPES[selected_tower_type]["cost"]
                            selected_tower_type = None
                # 点击已有炮塔查看信息
                for t in towers:
                    if math.sqrt((mx - t.x)**2 + (my - t.y)**2) < 22:
                        tower_info_visible = True
                        info_tower = t
                        break

        elif ev.type == pygame.KEYDOWN:
            if ev.key == pygame.K_r and game_over:
                reset_game()
            if ev.key == pygame.K_ESCAPE:
                selected_tower_type = None
                tower_info_visible = False

    # ===== 更新 =====
    if start_screen:
        draw_start_screen()
        pygame.display.flip()
        clock.tick(FPS)
        continue

    if not game_over:
        # 刷怪
        spawn_timer += 1
        if spawn_timer >= spawn_delay and len(enemies) < 8 and enemies_defeated < wave * enemies_per_wave:
            enemies.append(Enemy(path, wave))
            spawn_timer = 0

        # 敌人
        for e in enemies[:]:
            if e.update():
                health -= 1
                enemies.remove(e)
                if health <= 0:
                    game_over = True
            elif e.health <= 0:
                money += e.reward
                enemies_defeated += 1
                enemies.remove(e)
                explosions.append(Explosion(e.x, e.y))

        # 炮塔射击
        for t in towers:
            b = t.update(enemies)
            if b:
                bullets.append(b)

        # 子弹
        for b in bullets[:]:
            b.update()
            if not b.active and b in bullets:
                bullets.remove(b)

        # 爆炸
        for ex in explosions[:]:
            ex.update()
            if not ex.active:
                explosions.remove(ex)

        # 波次推进
        if enemies_defeated >= wave * enemies_per_wave and len(enemies) == 0:
            wave += 1
            money += 100
            if wave > 10:
                game_over = True
                game_won = True

    # ===== 绘制 =====
    screen.fill(BG)

    # 先画网格（在草地下面一层）
    grid = build_grid(towers)
    draw_grid(screen, grid, mouse_pos)

    draw_grass_bg()
    draw_path_line()
    draw_radish()

    # 选中炮塔时显示预览
    if selected_tower_type:
        mx, my = mouse_pos
        cell = point_to_cell(mx, my)
        cx, cy = cell_center(*cell)
        gstatus = grid.get(cell, "no")
        if gstatus == "ok":
            # 预览范围
            r = TOWER_TYPES[selected_tower_type]["range"]
            s2 = pygame.Surface((r*2, r*2), pygame.SRCALPHA)
            pygame.draw.circle(s2, (80, 200, 80, 50), (r, r), r, 0)
            pygame.draw.circle(s2, (80, 255, 80, 120), (r, r), r, 2)
            screen.blit(s2, (cx - r, cy - r))
            # 预览炮塔
            tc = TOWER_TYPES[selected_tower_type]["color"]
            pygame.draw.circle(screen, tc, (cx, cy), 20)
            pygame.draw.circle(screen, (max(0,tc[0]-30),max(0,tc[1]-30),max(0,tc[2]-30)), (cx, cy), 15)
            pygame.draw.rect(screen, (max(0,tc[0]-50),max(0,tc[1]-50),max(0,tc[2]-50)), (cx+5, cy-5, 22, 10))
        else:
            # 红色叉号提示不可放
            pygame.draw.line(screen, (255, 80, 80), (cx-15, cy-15), (cx+15, cy+15), 4)
            pygame.draw.line(screen, (255, 80, 80), (cx+15, cy-15), (cx-15, cy+15), 4)

    # 炮塔
    for t in towers:
        t.draw(screen)
        # 悬停显示范围
        if math.sqrt((mouse_pos[0]-t.x)**2 + (mouse_pos[1]-t.y)**2) < 25:
            s3 = pygame.Surface((t.range*2, t.range*2), pygame.SRCALPHA)
            pygame.draw.circle(s3, (80, 130, 255, 35), (t.range, t.range), t.range, 0)
            pygame.draw.circle(s3, (80, 160, 255, 100), (t.range, t.range), t.range, 1)
            screen.blit(s3, (t.x - t.range, t.y - t.range))

    # 子弹 / 敌人 / 爆炸
    for b in bullets:
        b.draw(screen)
    for e in enemies:
        e.draw(screen)
    for ex in explosions:
        ex.draw(screen)

    # UI
    draw_ui()
    draw_tower_info_panel()
    if game_over:
        draw_game_over_screen()

    pygame.display.flip()
    clock.tick(FPS)

pygame.quit()
sys.exit()
