import pygame
import random
import math
import sys

pygame.init()
WIDTH = 820
HEIGHT = 560
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Protect The Watermelon")

# Colors
SKY_TOP = (135, 206, 235)
SKY_BOTTOM = (180, 225, 245)
GRASS_DARK = (45, 120, 35)
GRASS_MAIN = (60, 140, 50)
GRASS_LIGHT = (90, 170, 70)
PATH_COLOR = (190, 160, 120)
PATH_DARK = (170, 135, 95)
PATH_EDGE = (140, 110, 70)
WATERMELON_DARK = (20, 90, 20)
WATERMELON_LIGHT = (40, 130, 40)
WATERMELON_RED = (255, 80, 80)
WATERMELON_SEED = (30, 20, 20)
TOWER_GUN = (70, 70, 70)
TOWER_GUN_LIGHT = (100, 100, 100)
TOWER_ICE = (70, 170, 210)
TOWER_ICE_LIGHT = (100, 200, 240)
ENEMY_NORMAL = (140, 30, 30)
ENEMY_FAST = (180, 110, 30)
BULLET = (255, 255, 80)
ICE_BULLET = (150, 220, 255)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GOLD = (255, 210, 0)
RED = (220, 30, 30)
GREEN = (40, 200, 40)
UI_BG = (20, 20, 20)

# Font
font_big = pygame.font.Font(None, 42)
font_mid = pygame.font.Font(None, 32)
font_small = pygame.font.Font(None, 26)

# Path
path = [
    (0, 80),
    (200, 80),
    (200, 220),
    (80, 220),
    (80, 360),
    (320, 360),
    (320, 160),
    (520, 160),
    (520, 420),
    (700, 420),
    (700, 280),
    (760, 280)
]
watermelon_pos = (760, 280)
PATH_THICK = 52

# Game Data
MAX_LEVEL = 5
current_level = 1
gold = 120
watermelon_hp = 10
wave = 1
max_wave = 6
game_over = False
level_clear = False
game_victory = False
spawn_timer = 0
wave_enemy_count = 5
spawned = 0
clock = pygame.time.Clock()

enemies = []
towers = []
bullets = []
grass_dots = []
clouds = []
particles = []

# Tower settings
TOWER_GUN_COST = 50
TOWER_ICE_COST = 75
selected_tower = 0

for _ in range(160):
    grass_dots.append({
        "x": random.randint(0, WIDTH),
        "y": random.randint(0, 460),
        "size": random.choice([1, 2, 3]),
        "color": random.choice([GRASS_LIGHT, (80, 160, 60), (100, 180, 80)])
    })

for _ in range(6):
    clouds.append({
        "x": random.randint(0, WIDTH),
        "y": random.randint(20, 90),
        "w": random.randint(80, 120),
        "speed": random.uniform(0.15, 0.35)
    })

class Particle:
    def __init__(self, x, y, color):
        self.x = x
        self.y = y
        self.vx = random.uniform(-2, 2)
        self.vy = random.uniform(-2, 2)
        self.life = 30
        self.color = color
    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.life -= 1
    def draw(self):
        pygame.draw.circle(screen, self.color, (int(self.x), int(self.y)), 2)

class Enemy:
    def __init__(self, speed_mult=1):
        self.x, self.y = path[0]
        self.target_idx = 1
        self.hp = 60 * speed_mult
        self.max_hp = self.hp
        self.speed = 1.1 * speed_mult
        self.slow_timer = 0
        self.anim_timer = 0
    def update(self):
        self.anim_timer += 1
        if self.slow_timer > 0:
            spd = self.speed * 0.4
            self.slow_timer -= 1
        else:
            spd = self.speed
        tx, ty = path[self.target_idx]
        dx = tx - self.x
        dy = ty - self.y
        dist = math.hypot(dx, dy)
        if dist < spd:
            self.x = tx
            self.y = ty
            self.target_idx += 1
        else:
            self.x += dx / dist * spd
            self.y += dy / dist * spd
    def draw(self):
        r = 16
        bounce = math.sin(self.anim_timer * 0.15) * 2
        if self.max_hp > 60:
            base_color = ENEMY_FAST
            shadow_color = (150, 90, 20)
        else:
            base_color = ENEMY_NORMAL
            shadow_color = (110, 20, 20)
        pygame.draw.circle(screen, shadow_color, (int(self.x + 2), int(self.y + 2 + bounce)), r)
        pygame.draw.circle(screen, base_color, (int(self.x), int(self.y + bounce)), r)
        pygame.draw.circle(screen, WHITE, (int(self.x - 5), int(self.y - 5 + bounce)), 5)
        pygame.draw.circle(screen, WHITE, (int(self.x + 5), int(self.y - 5 + bounce)), 5)
        pygame.draw.circle(screen, BLACK, (int(self.x - 4), int(self.y - 4 + bounce)), 2)
        pygame.draw.circle(screen, BLACK, (int(self.x + 6), int(self.y - 4 + bounce)), 2)
        bar_w = 28
        hp_ratio = self.hp / self.max_hp
        pygame.draw.rect(screen, BLACK, (self.x - bar_w // 2, self.y - 26, bar_w, 5))
        pygame.draw.rect(screen, GREEN if hp_ratio > 0.5 else RED,
                         (self.x - bar_w // 2, self.y - 26, bar_w * hp_ratio, 5))
        if self.slow_timer > 0:
            pygame.draw.circle(screen, ICE_BULLET, (int(self.x), int(self.y + bounce)), r + 2, 2)

class Tower:
    def __init__(self, x, y, tower_type):
        self.x = x
        self.y = y
        self.type = tower_type
        self.range = 130
        self.cooldown = 0
        self.shoot_anim = 0
        if tower_type == 1:
            self.damage = 12
            self.cd_max = 45
        else:
            self.damage = 6
            self.cd_max = 70
    def update(self):
        if self.shoot_anim > 0:
            self.shoot_anim -= 1
        self.cooldown -= 1
        if self.cooldown <= 0:
            target = None
            min_dist = self.range
            for e in enemies:
                d = math.hypot(e.x - self.x, e.y - self.y)
                if d < min_dist:
                    min_dist = d
                    target = e
            if target is not None:
                bullets.append({
                    "x": self.x, "y": self.y,
                    "tx": target.x, "ty": target.y,
                    "target": target, "speed": 6,
                    "dmg": self.damage, "is_ice": self.type == 2
                })
                self.cooldown = self.cd_max
                self.shoot_anim = 8
                for _ in range(5):
                    particles.append(Particle(self.x, self.y, ICE_BULLET if self.type == 2 else BULLET))
    def draw(self):
        scale = 1 + self.shoot_anim * 0.015
        r = int(18 * scale)
        pygame.draw.circle(screen, (0, 0, 0), (self.x + 2, self.y + 2), r)
        if self.type == 1:
            pygame.draw.circle(screen, TOWER_GUN_LIGHT, (self.x, self.y), r)
            pygame.draw.circle(screen, TOWER_GUN, (self.x, self.y), r, 3)
            pygame.draw.rect(screen, BLACK, (self.x - 4, self.y - 24, 8, 14), border_radius=2)
        else:
            pygame.draw.circle(screen, TOWER_ICE_LIGHT, (self.x, self.y), r)
            pygame.draw.circle(screen, TOWER_ICE, (self.x, self.y), r, 3)
            pygame.draw.circle(screen, WHITE, (self.x, self.y), r - 6, 2)

def reset_level():
    global gold, watermelon_hp, wave, game_over, level_clear, spawn_timer, wave_enemy_count, spawned
    global enemies, towers, bullets, particles
    gold = 120
    watermelon_hp = 10
    wave = 1
    game_over = False
    level_clear = False
    spawn_timer = 0
    wave_enemy_count = 5 + (current_level - 1) * 2
    spawned = 0
    enemies.clear()
    towers.clear()
    bullets.clear()
    particles.clear()

def reset_all_game():
    global current_level, game_victory
    current_level = 1
    game_victory = False
    reset_level()

reset_all_game()

def draw_background():
    for y in range(100):
        ratio = y / 100
        r = int(SKY_TOP[0] * (1 - ratio) + SKY_BOTTOM[0] * ratio)
        g = int(SKY_TOP[1] * (1 - ratio) + SKY_BOTTOM[1] * ratio)
        b = int(SKY_TOP[2] * (1 - ratio) + SKY_BOTTOM[2] * ratio)
        pygame.draw.line(screen, (r, g, b), (0, y), (WIDTH, y))
    for cloud in clouds:
        cloud["x"] += cloud["speed"]
        if cloud["x"] > WIDTH + 120:
            cloud["x"] = -120
        x = int(cloud["x"])
        y = int(cloud["y"])
        w = int(cloud["w"])
        pygame.draw.ellipse(screen, WHITE, (x, y, w, 35))
        pygame.draw.ellipse(screen, WHITE, (x + w // 4, y - 10, w // 2, 30))
        pygame.draw.ellipse(screen, (240, 240, 255), (x + 10, y + 5, w // 3, 20))
    screen.fill(GRASS_MAIN, (0, 100, WIDTH, HEIGHT - 100))
    for dot in grass_dots:
        pygame.draw.circle(screen, dot["color"], (dot["x"], dot["y"] + 100), dot["size"])
    for i in range(len(path) - 1):
        x1, y1 = path[i]
        x2, y2 = path[i + 1]
        pygame.draw.line(screen, PATH_EDGE, (x1, y1 + 3), (x2, y2 + 3), PATH_THICK + 4)
        pygame.draw.line(screen, PATH_DARK, (x1, y1), (x2, y2), PATH_THICK)
        pygame.draw.line(screen, PATH_COLOR, (x1, y1), (x2, y2), PATH_THICK - 6)

def draw_watermelon(x, y):
    pygame.draw.circle(screen, BLACK, (x + 3, y + 3), 26)
    pygame.draw.circle(screen, WATERMELON_DARK, (x, y), 24)
    pygame.draw.circle(screen, WATERMELON_LIGHT, (x - 4, y - 4), 21)
    pygame.draw.circle(screen, WATERMELON_RED, (x, y), 12)
    for seed in [(-4, -4), (4, -3), (-2, 4), (5, 3), (0, 0)]:
        pygame.draw.circle(screen, WATERMELON_SEED, (x + seed[0], y + seed[1]), 2)

running = True
while running:
    dt = clock.tick(60)
    mx, my = pygame.mouse.get_pos()
    draw_background()

    draw_watermelon(watermelon_pos[0], watermelon_pos[1])

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_r:
                reset_all_game()
        if event.type == pygame.MOUSEBUTTONDOWN and not game_over and not level_clear and not game_victory:
            if event.button == 1:
                if 10 <= mx <= 130 and 480 <= my <= 530:
                    selected_tower = 1
                elif 145 <= mx <= 265 and 480 <= my <= 530:
                    selected_tower = 2
                else:
                    if selected_tower != 0:
                        cost = TOWER_GUN_COST if selected_tower == 1 else TOWER_ICE_COST
                        if gold >= cost:
                            towers.append(Tower(mx, my, selected_tower))
                            gold -= cost

    if not game_over and not level_clear and not game_victory:
        spawn_timer += dt
        if spawn_timer > 1100 and spawned < wave_enemy_count:
            spawn_timer = 0
            fast_mult = 1.6 if random.random() < 0.35 else 1
            enemies.append(Enemy(fast_mult))
            spawned += 1

        for e in enemies[:]:
            e.update()
            if e.target_idx >= len(path):
                watermelon_hp -= 1
                enemies.remove(e)
                if watermelon_hp <= 0:
                    game_over = True
            elif e.hp <= 0:
                enemies.remove(e)
                gold += 12
                for _ in range(8):
                    particles.append(Particle(e.x, e.y, random.choice([BULLET, WHITE])))

        if spawned >= wave_enemy_count and len(enemies) == 0:
            wave += 1
            spawned = 0
            spawn_timer = -1800
            if wave > max_wave:
                level_clear = True

        for t in towers:
            t.update()

        for b in bullets[:]:
            if b["target"] in enemies:
                tx = b["target"].x
                ty = b["target"].y
            else:
                tx = b["tx"]
                ty = b["ty"]
            dx = tx - b["x"]
            dy = ty - b["y"]
            dist = math.hypot(dx, dy)
            if dist < b["speed"]:
                if b["target"] in enemies:
                    b["target"].hp -= b["dmg"]
                    if b["is_ice"]:
                        b["target"].slow_timer = 90
                bullets.remove(b)
            else:
                b["x"] += dx / dist * b["speed"]
                b["y"] += dy / dist * b["speed"]

        for p in particles[:]:
            p.update()
            if p.life <= 0:
                particles.remove(p)

    if level_clear:
        current_level += 1
        if current_level > MAX_LEVEL:
            game_victory = True
        else:
            reset_level()
            level_clear = False

    for p in particles:
        p.draw()
    for t in towers:
        t.draw()
    for e in enemies:
        e.draw()
    for b in bullets:
        if b["is_ice"]:
            pygame.draw.circle(screen, ICE_BULLET, (int(b["x"]), int(b["y"])), 5)
            pygame.draw.circle(screen, WHITE, (int(b["x"]), int(b["y"])), 5, 1)
        else:
            pygame.draw.circle(screen, BULLET, (int(b["x"]), int(b["y"])), 5)
            pygame.draw.circle(screen, WHITE, (int(b["x"]), int(b["y"])), 5, 1)

    # ======修复这里：矩形参数改成元组======
    pygame.draw.rect(screen, UI_BG, (0, 460, WIDTH, 100))
    level_text = font_mid.render(f"Level {current_level}/{MAX_LEVEL}", True, WHITE)
    gold_text = font_mid.render(f"Gold: {gold}", True, GOLD)
    hp_text = font_mid.render(f"Watermelon HP: {watermelon_hp}", True, RED)
    wave_text = font_mid.render(f"Wave: {wave}/{max_wave}", True, WHITE)
    screen.blit(level_text, (12, 465))
    screen.blit(gold_text, (130, 465))
    screen.blit(hp_text, (260, 465))
    screen.blit(wave_text, (400, 465))

    col1 = TOWER_GUN if selected_tower == 1 else (50, 50, 50)
    pygame.draw.rect(screen, col1, (10, 480, 120, 50), border_radius=6)
    pygame.draw.rect(screen, WHITE, (10, 480, 120, 50), 2, border_radius=6)
    t1_txt = font_small.render(f"Gun ${TOWER_GUN_COST}", True, WHITE)
    screen.blit(t1_txt, t1_txt.get_rect(center=(70, 505)))

    col2 = TOWER_ICE if selected_tower == 2 else (50, 50, 50)
    pygame.draw.rect(screen, col2, (145, 480, 120, 50), border_radius=6)
    pygame.draw.rect(screen, WHITE, (145, 480, 120, 50), 2, border_radius=6)
    t2_txt = font_small.render(f"Ice ${TOWER_ICE_COST}", True, WHITE)
    screen.blit(t2_txt, t2_txt.get_rect(center=(205, 505)))

    hint = font_small.render("Click tower button, then click grass to place", True, WHITE)
    screen.blit(hint, (280, 490))

    overlay = pygame.Surface((WIDTH, HEIGHT))
    overlay.set_alpha(140)
    if game_over:
        overlay.fill(BLACK)
        screen.blit(overlay, (0,0))
        msg1 = font_big.render("Game Over!", True, RED)
        msg2 = font_mid.render("Press R to restart all game", True, WHITE)
        screen.blit(msg1, msg1.get_rect(center=(WIDTH//2, HEIGHT//2-40)))
        screen.blit(msg2, msg2.get_rect(center=(WIDTH//2, HEIGHT//2+20)))
    elif game_victory:
        overlay.fill(BLACK)
        screen.blit(overlay, (0,0))
        msg1 = font_big.render("ALL LEVELS CLEARED! YOU WIN!", True, GOLD)
        msg2 = font_mid.render("Press R to play again", True, WHITE)
        screen.blit(msg1, msg1.get_rect(center=(WIDTH//2, HEIGHT//2-40)))
        screen.blit(msg2, msg2.get_rect(center=(WIDTH//2, HEIGHT//2+20)))

    pygame.display.update()

pygame.quit()
sys.exit()