import pygame
import sys
import math
import random

# 初始化pygame
pygame.init()

# 屏幕设置
WIDTH, HEIGHT = 980, 720
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Radish Defense")

# 颜色定义
BACKGROUND = (30, 60, 30)
PATH_COLOR = (180, 160, 100)
GRASS_COLOR = (70, 140, 70)
BULLET_COLOR = (255, 200, 0)
ENEMY_COLOR = (220, 80, 80)
HEALTH_BAR_GREEN = (80, 220, 80)
HEALTH_BAR_RED = (220, 80, 80)
UI_BG = (40, 40, 60)
UI_TEXT = (220, 220, 255)
BUTTON_COLOR = (70, 100, 150)
BUTTON_HOVER = (90, 130, 190)
RADISH_COLOR = (240, 120, 160)
RADISH_LEAVES = (100, 220, 100)
SLOW_COLOR = (100, 200, 255)
EXPLOSION_COLOR = (255, 150, 0)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

# 游戏参数
FPS = 60
clock = pygame.time.Clock()

# 使用 Font(None, size)
font = pygame.font.Font(None, 24)
big_font = pygame.font.Font(None, 48)
title_font = pygame.font.Font(None, 72)

# 路径定义
path = [
    (50, 358), (228, 362), (239, 218), (411, 214),
    (406, 401), (593, 404), (624, 292), (811, 323), (862, 538)
]

# 游戏状态
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": 145, "damage": 1, "cooldown": 30, "color": (180, 182, 200)},
    "sniper": {"name": "Sniper", "cost": 115, "range": 295, "damage": 3, "cooldown": 60, "color": (100, 195, 105)},
    "machine": {"name": "Machine", "cost": 95, "range": 120, "damage": 0.5, "cooldown": 10, "color": (195, 105, 108)},
    "cannon": {"name": "Cannon", "cost": 145, "range": 130, "damage": 2, "cooldown": 40, "color": (195, 150, 52)},
    "freeze": {"name": "Freeze", "cost": 128, "range": 140, "damage": 0.2, "cooldown": 50, "color": (100, 180, 255)}
}


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_color = (max(0, self.color[0]-50), max(0, self.color[1]-50), max(0, self.color[2]-50))

    def draw(self, surface):
        pygame.draw.circle(surface, 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(surface, darker, (self.x, self.y), 15)

        if self.type == "basic":
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y-5, 25, 10))
        elif self.type == "sniper":
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y-3, 35, 6))
            pygame.draw.circle(surface, (50, 50, 50), (self.x+40, self.y), 5)
        elif self.type == "machine":
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y-8, 25, 6))
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y+2, 25, 6))
        elif self.type == "cannon":
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y-7, 30, 14))
        elif self.type == "freeze":
            pygame.draw.rect(surface, self.gun_color, (self.x+5, self.y-5, 25, 10))
            pygame.draw.circle(surface, (200, 230, 255), (self.x+15, self.y), 8, 2)

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


class Bullet:
    def __init__(self, x, y, target, damage, bullet_type):
        self.x = x
        self.y = y
        self.target = target
        self.speed = 8
        self.damage = damage
        self.active = True
        self.type = bullet_type
        self.effect = None
        self.size = 5

        if bullet_type == "freeze":
            self.color = (100, 200, 255)
            self.effect = "slow"
            self.effect_duration = 180
        elif bullet_type == "cannon":
            self.color = (255, 150, 0)
            self.size = 8
        else:
            self.color = BULLET_COLOR

    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 *= 0.5
            self.active = False
            return True
        self.x += dx / dist * self.speed
        self.y += dy / dist * self.speed
        return False

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


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
        self.type = random.choice(["normal", "fast", "tank"])
        if self.type == "fast":
            self.speed *= 1.8
            self.health *= 0.7
            self.reward = int(self.reward * 1.3)
            self.color = (220, 120, 120)
        elif 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.color = ENEMY_COLOR

    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
        target_x, target_y = self.path[self.path_index]
        dx = target_x - self.x
        dy = target_y - 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
            target_x, target_y = self.path[self.path_index]
            dx = target_x - self.x
            dy = target_y - 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, surface):
        pygame.draw.circle(surface, 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(surface, darker, (int(self.x), int(self.y)), 10)
            pygame.draw.line(surface, (255, 255, 200), (self.x, self.y-15), (self.x, self.y-25), 2)
        elif self.type == "tank":
            dark2 = (max(0, self.color[0]-50), max(0, self.color[1]-50), max(0, self.color[2]-50))
            pygame.draw.rect(surface, dark2, (self.x-12, self.y-5, 24, 10))
            pygame.draw.rect(surface, dark2, (self.x-5, self.y-12, 10, 24))
        else:
            pygame.draw.circle(surface, darker, (int(self.x), int(self.y)), 10)
        bar_width = 40
        bar_height = 5
        health_ratio = self.health / self.max_health
        pygame.draw.rect(surface, HEALTH_BAR_RED, (self.x - bar_width//2, self.y - 30, bar_width, bar_height))
        pygame.draw.rect(surface, HEALTH_BAR_GREEN, (self.x - bar_width//2, self.y - 30, int(bar_width * health_ratio), bar_height))
        if self.slow_duration > 0:
            pygame.draw.circle(surface, SLOW_COLOR, (int(self.x), int(self.y)), 18, 2)


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

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

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


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


def draw_path():
    for i in range(len(path) - 1):
        pygame.draw.line(screen, PATH_COLOR, path[i], path[i+1], 40)
    pygame.draw.circle(screen, (120, 90, 60), path[0], 20)
    pygame.draw.circle(screen, (140, 110, 80), path[0], 15)
    pygame.draw.circle(screen, (160, 130, 100), path[0], 10)


def draw_grass():
    for i in range(0, WIDTH, 40):
        for j in range(0, HEIGHT, 40):
            if random.random() > 0.3:
                pygame.draw.line(screen, GRASS_COLOR, (i, j), (i, j-15), 2)


def draw_radish():
    x, y = path[-1]
    pygame.draw.ellipse(screen, RADISH_LEAVES, (x-25, y-70, 50, 40))
    leaf_dark = (max(0, RADISH_LEAVES[0]-20), max(0, RADISH_LEAVES[1]-20), max(0, RADISH_LEAVES[2]-20))
    pygame.draw.ellipse(screen, leaf_dark, (x-30, y-60, 30, 30))
    pygame.draw.ellipse(screen, leaf_dark, (x, y-60, 30, 30))
    pygame.draw.circle(screen, RADISH_COLOR, (x, y), 30)
    radish_dark = (max(0, RADISH_COLOR[0]-40), max(0, RADISH_COLOR[1]-40), max(0, RADISH_COLOR[2]-40))
    pygame.draw.ellipse(screen, radish_dark, (x-20, y-15, 40, 30))
    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, (200, 80, 120), (x-10, y, 20, 15), 0, math.pi, 3)
    pygame.draw.rect(screen, HEALTH_BAR_RED, (x-30, y-95, 60, 10))
    pygame.draw.rect(screen, HEALTH_BAR_GREEN, (x-30, y-95, int(60 * (health/10)), 10))


def draw_ui():
    pygame.draw.rect(screen, UI_BG, (0, 0, WIDTH, 50))
    money_text = font.render(f"Gold: {money}", True, UI_TEXT)
    screen.blit(money_text, (20, 15))
    health_text = font.render(f"Health: {health}", True, UI_TEXT)
    screen.blit(health_text, (150, 15))
    wave_text = font.render(f"Wave: {wave}/10", True, UI_TEXT)
    screen.blit(wave_text, (280, 15))
    defeated_text = font.render(f"Killed: {enemies_defeated}/{wave * enemies_per_wave}", True, UI_TEXT)
    screen.blit(defeated_text, (415, 15))

    ui_y = HEIGHT - 155
    pygame.draw.rect(screen, (50, 50, 70), (0, ui_y, WIDTH, 155))
    pygame.draw.line(screen, (80, 80, 100), (0, ui_y), (WIDTH, ui_y), 3)
    towers_text = font.render("Select Tower Type:", True, UI_TEXT)
    screen.blit(towers_text, (20, ui_y + 15))

    button_width = 150
    button_height = 100
    spacing = 20
    start_x = 20

    for i, (tower_id, tower_data) in enumerate(TOWER_TYPES.items()):
        x = start_x + i * (button_width + spacing)
        y = ui_y + 40
        button_rect = pygame.Rect(x, y, button_width, button_height)
        if selected_tower_type == tower_id:
            color = (max(0, tower_data["color"][0]//2), max(0, tower_data["color"][1]//2), max(0, tower_data["color"][2]//2))
        elif button_rect.collidepoint(pygame.mouse.get_pos()):
            color = (max(0, int(tower_data["color"][0]/1.2)), max(0, int(tower_data["color"][1]/1.2)), max(0, int(tower_data["color"][2]/1.2)))
        else:
            color = tower_data["color"]
        pygame.draw.rect(screen, color, button_rect)
        border = (max(0, color[0]//2), max(0, color[1]//2), max(0, color[2]//2))
        pygame.draw.rect(screen, border, button_rect, 3)
        name_text = font.render(tower_data["name"], True, WHITE)
        screen.blit(name_text, (x + button_width//2 - name_text.get_width()//2, y + 10))
        cost_text = font.render(f"Cost: {tower_data['cost']}", True, (255, 255, 200))
        screen.blit(cost_text, (x + button_width//2 - cost_text.get_width()//2, y + 40))
        icon_color = (max(0, color[0]-50), max(0, color[1]-50), max(0, color[2]-50))
        pygame.draw.circle(screen, icon_color, (x + button_width//2, y + 75), 15)
        if tower_id == "sniper":
            pygame.draw.rect(screen, (50, 50, 50), (x + button_width//2 + 5, y + 75 - 3, 20, 6))
        elif tower_id == "machine":
            pygame.draw.rect(screen, (50, 50, 50), (x + button_width//2 + 5, y + 75 - 8, 20, 6))
            pygame.draw.rect(screen, (50, 50, 50), (x + button_width//2 + 5, y + 75 + 2, 20, 6))
        elif tower_id == "cannon":
            pygame.draw.rect(screen, (50, 50, 50), (x + button_width//2 + 5, y + 75 - 7, 25, 14))
        elif tower_id == "freeze":
            pygame.draw.circle(screen, (200, 230, 255), (x + button_width//2, y + 75), 10, 2)


def draw_tower_info():
    if not tower_info_visible or not info_tower:
        return
    info_surface = pygame.Surface((300, 200), pygame.SRCALPHA)
    info_surface.fill((30, 30, 50, 220))
    pygame.draw.rect(info_surface, (100, 100, 150), (0, 0, 300, 200), 3)
    title = font.render(f"{TOWER_TYPES[info_tower.type]['name']} Info", True, (255, 255, 200))
    info_surface.blit(title, (150 - title.get_width()//2, 15))
    damage_text = font.render(f"Damage: {info_tower.damage}", True, UI_TEXT)
    info_surface.blit(damage_text, (30, 50))
    range_text = font.render(f"Range: {info_tower.range}", True, UI_TEXT)
    info_surface.blit(range_text, (30, 80))
    speed_text = font.render(f"Speed: {60/(info_tower.cooldown_max/FPS):.1f}/s", True, UI_TEXT)
    info_surface.blit(speed_text, (30, 110))
    if info_tower.type == "freeze":
        effect_text = font.render("Effect: Slow 50%", True, SLOW_COLOR)
        info_surface.blit(effect_text, (30, 140))
    pygame.draw.rect(info_surface, (180, 80, 80), (250, 10, 40, 20))
    close_text = font.render("X", True, WHITE)
    info_surface.blit(close_text, (265, 12))
    screen.blit(info_surface, (WIDTH//2 - 150, HEIGHT//2 - 100))


def draw_game_over():
    overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
    overlay.fill((0, 0, 0, 180))
    screen.blit(overlay, (0, 0))
    if game_won:
        text = big_font.render("Victory! Radish is safe!", True, (100, 255, 100))
    else:
        text = big_font.render("Game Over! Radish was eaten!", True, (255, 100, 100))
    text_rect = text.get_rect(center=(WIDTH//2, HEIGHT//2 - 50))
    screen.blit(text, text_rect)
    restart_text = font.render("Press R to Restart", True, (200, 200, 255))
    restart_rect = restart_text.get_rect(center=(WIDTH//2, HEIGHT//2 + 20))
    screen.blit(restart_text, restart_rect)
    score_text = font.render(f"Final Score: {enemies_defeated * 10 + money}", True, (255, 255, 150))
    score_rect = score_text.get_rect(center=(WIDTH//2, HEIGHT//2 + 60))
    screen.blit(score_text, score_rect)


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


def draw_start_screen():
    screen.fill((20, 40, 20))
    title_text = title_font.render("Radish Defense", True, (100, 220, 100))
    screen.blit(title_text, (WIDTH//2 - title_text.get_width()//2, 150))
    pygame.draw.circle(screen, RADISH_COLOR, (WIDTH//2, 290), 60)
    radish_dark = (max(0, RADISH_COLOR[0]-40), max(0, RADISH_COLOR[1]-40), max(0, RADISH_COLOR[2]-40))
    pygame.draw.ellipse(screen, radish_dark, (WIDTH//2-40, 375, 80, 50))
    pygame.draw.circle(screen, WHITE, (WIDTH//2-20, 395), 15)
    pygame.draw.circle(screen, WHITE, (WIDTH//2+20, 403), 15)
    pygame.draw.circle(screen, BLACK, (WIDTH//2-20, 397), 7)
    pygame.draw.circle(screen, BLACK, (WIDTH//2+20, 407), 7)
    pygame.draw.arc(screen, (200, 80, 120), (WIDTH//2-20, 385, 40, 30), 0, math.pi, 5)
    pygame.draw.ellipse(screen, RADISH_LEAVES, (WIDTH//2-40, 289, 80, 60))
    leaf_dark = (max(0, RADISH_LEAVES[0]-20), max(0, RADISH_LEAVES[1]-20), max(0, RADISH_LEAVES[2]-20))
    pygame.draw.ellipse(screen, leaf_dark, (WIDTH//2-60, 318, 40, 40))
    pygame.draw.ellipse(screen, leaf_dark, (WIDTH//2+20, 348, 40, 40))
    btn_rect = pygame.Rect(WIDTH//2-100, 495, 200, 60)
    pygame.draw.rect(screen, BUTTON_COLOR, btn_rect)
    pygame.draw.rect(screen, BUTTON_HOVER, btn_rect, 3)
    start_text = big_font.render("Start Game", True, WHITE)
    screen.blit(start_text, (WIDTH//2 - start_text.get_width()//2, 508))


running = True
start_screen = True

while running:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            if start_screen:
                btn_rect = pygame.Rect(WIDTH//2-100, 493, 200, 60)
                if btn_rect.collidepoint(event.pos):
                    start_screen = False
                continue
            if game_over:
                continue
            x, y = event.pos
            if tower_info_visible:
                if WIDTH//2+100 <= x <= WIDTH//2+140 and HEIGHT//2-100 <= y <= HEIGHT//2-80:
                    tower_info_visible = False
                continue
            ui_y = HEIGHT - 155
            if y >= ui_y:
                button_width = 150
                spacing = 20
                start_x = 20
                for i, tower_id in enumerate(TOWER_TYPES.keys()):
                    btn_x = start_x + i * (button_width + spacing)
                    btn_y = ui_y + 40
                    if btn_x <= x <= btn_x + button_width and btn_y <= y <= btn_y + 100:
                        if money >= TOWER_TYPES[tower_id]["cost"]:
                            selected_tower_type = tower_id
                        break
            elif selected_tower_type and y < ui_y:
                on_path = False
                for i in range(len(path) - 1):
                    p1 = path[i]
                    p2 = path[i+1]
                    dist = abs((p2[0]-p1[0])*(p1[1]-y) - (p1[0]-x)*(p2[1]-p1[1])) / math.sqrt((p2[0]-p1[0])**2 + (p2[1]-p1[1])**2)
                    if dist < 30:
                        on_path = True
                        break
                radish_dist = math.sqrt((x - path[-1][0])**2 + (y - path[-1][1])**2)
                if radish_dist < 50:
                    on_path = True
                too_close = any(math.sqrt((x - t.x)**2 + (y - t.y)**2) < 50 for t in towers)
                if not on_path and not too_close and money >= TOWER_TYPES[selected_tower_type]["cost"]:
                    towers.append(Tower(x, y, selected_tower_type))
                    money -= TOWER_TYPES[selected_tower_type]["cost"]
                    selected_tower_type = None
                for tower in towers:
                    if math.sqrt((x - tower.x)**2 + (y - tower.y)**2) < 25:
                        tower_info_visible = True
                        info_tower = tower
                        break
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_r and game_over:
                reset_game()
            if event.key == pygame.K_ESCAPE:
                selected_tower_type = None

    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 enemy in enemies[:]:
            reached_end = enemy.update()
            if reached_end:
                health -= 1
                enemies.remove(enemy)
                if health <= 0:
                    game_over = True
            elif enemy.health <= 0:
                money += enemy.reward
                enemies_defeated += 1
                enemies.remove(enemy)
                explosions.append(Explosion(enemy.x, enemy.y))
        for tower in towers:
            bullet = tower.update(enemies)
            if bullet:
                bullets.append(bullet)
        for bullet in bullets[:]:
            hit = bullet.update()
            if hit or not bullet.active:
                if bullet in bullets:
                    bullets.remove(bullet)
        for explosion in explosions[:]:
            explosion.update()
            if not explosion.active:
                explosions.remove(explosion)
        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(BACKGROUND)
    draw_grass()
    draw_path()
    draw_radish()
    if selected_tower_type:
        mouse_pos = pygame.mouse.get_pos()
        pygame.draw.circle(screen, (100, 100, 200), mouse_pos, TOWER_TYPES[selected_tower_type]["range"], 2)
        pygame.draw.circle(screen, TOWER_TYPES[selected_tower_type]["color"], mouse_pos, 20)
    for tower in towers:
        tower.draw(screen)
        mouse_pos = pygame.mouse.get_pos()
        if math.sqrt((mouse_pos[0] - tower.x)**2 + (mouse_pos[1] - tower.y)**2) < 25:
            pygame.draw.circle(screen, (100, 100, 200), (tower.x, tower.y), tower.range, 1)
    for bullet in bullets:
        bullet.draw(screen)
    for enemy in enemies:
        enemy.draw(screen)
    for explosion in explosions:
        explosion.draw(screen)
    draw_ui()
    draw_tower_info()
    if game_over:
        draw_game_over()
    pygame.display.flip()
    clock.tick(FPS)

pygame.quit()
sys.exit()