import pygame
import random
import math

# --- 初始化 ---
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("植物大战僵尸 - 进阶版")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (120, 180, 80)
CARD_BG = (200, 180, 140)
SUN_COLOR = (255, 220, 50)
PEASHOOTER_COLOR = (50, 180, 50)
WALLNUT_COLOR = (180, 120, 60)
SUNFLOWER_COLOR = (255, 200, 50)
CHERRY_COLOR = (200, 30, 30)
ZOMBIE_COLOR = (100, 120, 100)
BULLET_COLOR = (100, 200, 100)
UI_COLOR = (255, 255, 255)

FONT = pygame.font.SysFont("arial", 24)
BIG_FONT = pygame.font.SysFont("arial", 60)

# 网格参数
GRID_X, GRID_Y = 100, 100
CELL_W, CELL_H = 80, 100
ROWS, COLS = 5, 9

# --- 游戏对象 ---

class Sun:
    def __init__(self, x, y, target_y=None):
        self.x = x
        self.y = y
        self.target_y = target_y if target_y else y
        self.radius = 20
        self.alive = True
        self.falling = target_y is not None
        self.life = 400

    def update(self):
        if self.falling and self.y < self.target_y:
            self.y += 1
        else:
            self.falling = False
        self.life -= 1
        if self.life <= 0: self.alive = False

    def draw(self, surface):
        pygame.draw.circle(surface, SUN_COLOR, (int(self.x), int(self.y)), self.radius)
        pygame.draw.circle(surface, (255, 255, 200), (int(self.x), int(self.y)), self.radius - 5)

class Bullet:
    def __init__(self, x, y, row):
        self.x = x
        self.y = y
        self.row = row
        self.speed = 8
        self.alive = True

    def update(self):
        self.x += self.speed
        if self.x > WIDTH: self.alive = False

    def draw(self, surface):
        pygame.draw.circle(surface, BULLET_COLOR, (int(self.x), int(self.y)), 8)

class Plant:
    def __init__(self, col, row, ptype):
        self.col = col
        self.row = row
        self.x = GRID_X + col * CELL_W + CELL_W // 2
        self.y = GRID_Y + row * CELL_H + CELL_H // 2
        self.type = ptype
        self.hp = 150 if ptype == 'wallnut' else 50
        self.max_hp = self.hp
        self.cooldown = 0
        self.alive = True
        self.sun_timer = 0 if ptype == 'sunflower' else -1
        self.explode_timer = 60 if ptype == 'cherry' else -1

    def update(self, zombies, bullets, suns):
        if not self.alive: return

        if self.type == 'peashooter':
            has_zombie = any(z.row == self.row and z.x > self.x for z in zombies)
            if has_zombie and self.cooldown <= 0:
                bullets.append(Bullet(self.x + 20, self.y, self.row))
                self.cooldown = 60

        elif self.type == 'sunflower':
            self.sun_timer += 1
            if self.sun_timer >= 600:  # 10秒产一个
                self.sun_timer = 0
                suns.append(Sun(self.x, self.y - 20, self.y + 30))

        elif self.type == 'cherry':
            self.explode_timer -= 1
            if self.explode_timer <= 0:
                # 爆炸逻辑：清除周围 3x3
                for z in zombies:
                    if z.alive and abs(z.row - self.row) <= 1:
                        z_col = (z.x - GRID_X) // CELL_W
                        if abs(z_col - self.col) <= 1:
                            z.hp = 0
                            z.alive = False
                self.alive = False

        if self.cooldown > 0: self.cooldown -= 1

    def draw(self, surface):
        color_map = {
            'peashooter': PEASHOOTER_COLOR,
            'wallnut': WALLNUT_COLOR,
            'sunflower': SUNFLOWER_COLOR,
            'cherry': CHERRY_COLOR
        }
        color = color_map.get(self.type, (255,255,255))
        
        # 坚果被啃变色
        if self.type == 'wallnut' and self.hp < self.max_hp * 0.5:
            color = (120, 80, 40)

        pygame.draw.circle(surface, color, (self.x, self.y), 30)
        
        # 简单装饰
        if self.type == 'peashooter':
            pygame.draw.rect(surface, (0, 100, 0), (self.x + 10, self.y - 5, 20, 10))
        elif self.type == 'sunflower':
            pygame.draw.circle(surface, (100, 50, 0), (self.x, self.y), 10)
        elif self.type == 'cherry':
            pygame.draw.circle(surface, (0, 100, 0), (self.x - 10, self.y - 25), 5)
            pygame.draw.circle(surface, (0, 100, 0), (self.x + 10, self.y - 25), 5)

class Zombie:
    def __init__(self, row):
        self.row = row
        self.x = WIDTH + random.randint(0, 100)
        self.y = GRID_Y + row * CELL_H + CELL_H // 2
        self.speed = 0.5
        self.hp = 100
        self.alive = True
        self.attacking = False

    def update(self, plants):
        if not self.alive: return
        
        self.attacking = False
        for p in plants:
            if p.alive and p.row == self.row and abs(self.x - p.x) < 40:
                self.attacking = True
                p.hp -= 0.5
                if p.hp <= 0: p.alive = False
                break
        
        if not self.attacking:
            self.x -= self.speed

    def draw(self, surface):
        pygame.draw.rect(surface, ZOMBIE_COLOR, (int(self.x) - 20, int(self.y) - 30, 40, 60))
        pygame.draw.circle(surface, (150, 180, 150), (int(self.x), int(self.y) - 30), 15)
        if self.attacking:
            pygame.draw.line(surface, (255,0,0), (self.x-20, self.y-10), (self.x-30, self.y-10), 2)

# --- 主程序 ---

def main():
    sun_count = 150
    plants = []
    zombies = []
    bullets = []
    suns = []
    
    selected_plant = None
    spawn_timer = 0
    sun_timer = 0
    game_over = False

    running = True
    while running:
        clock.tick(60)
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT: running = False
            if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
                mx, my = pygame.mouse.get_pos()
                
                # 卡片点击
                if 50 <= mx <= 150 and 20 <= my <= 70: selected_plant = 'peashooter'
                elif 160 <= mx <= 260 and 20 <= my <= 70: selected_plant = 'wallnut'
                elif 270 <= mx <= 370 and 20 <= my <= 70: selected_plant = 'sunflower'
                elif 380 <= mx <= 480 and 20 <= my <= 70: selected_plant = 'cherry'
                
                # 阳光点击
                for s in suns:
                    if s.alive and math.hypot(mx - s.x, my - s.y) < 30:
                        sun_count += 25
                        s.alive = False
                
                # 种植
                    col = (mx - GRID_X) // CELL_W
                    row = (my - GRID_Y) // CELL_H
                    if 0 <= col < COLS and 0 <= row < ROWS:
                        if not any(p.col == col and p.row == row and p.alive for p in plants):
                            costs = {'peashooter': 100, 'wallnut': 50, 'sunflower': 50, 'cherry': 150}
                            cost = costs.get(selected_plant, 999)
                            if sun_count >= cost:
                                sun_count -= cost
                                plants.append(Plant(col, row, selected_plant))
                                selected_plant = None

        if not game_over:
            spawn_timer += 1
            if spawn_timer > 300:
                spawn_timer = 0
                zombies.append(Zombie(random.randint(0, 4)))

            sun_timer += 1
            if sun_timer > 500:
                sun_timer = 0
                suns.append(Sun(random.randint(100, 800), -20, random.randint(150, 500)))

            for p in plants: p.update(zombies, bullets, suns)
            for z in zombies: z.update(plants)
            for b in bullets: b.update()
            for s in suns: s.update()

            for b in bullets[:]:
                if not b.alive: 
                    bullets.remove(b)
                    continue
                for z in zombies:
                    if z.alive and z.row == b.row and abs(b.x - z.x) < 30:
                        z.hp -= 20
                        b.alive = False
                        if z.hp <= 0: z.alive = False
                        break
            
            plants = [p for p in plants if p.alive]
            zombies = [z for z in zombies if z.alive]
            suns = [s for s in suns if s.alive]

            if any(z.x < GRID_X - 50 for z in zombies):
                game_over = True

        # --- 绘图 ---
        screen.fill(BG_COLOR)
        
        for r in range(ROWS + 1):
            pygame.draw.line(screen, (100, 160, 60), (GRID_X, GRID_Y + r * CELL_H), (GRID_X + COLS * CELL_W, GRID_Y + r * CELL_H))
        for c in range(COLS + 1):
            pygame.draw.line(screen, (100, 160, 60), (GRID_X + c * CELL_W, GRID_Y), (GRID_X + c * CELL_W, GRID_Y + ROWS * CELL_H))

        # 卡片 UI
        cards = [("Pea 100", 50), ("Nut 50", 160), ("Sun 50", 270), ("Bomb 150", 380)]
        for txt, x in cards:
            pygame.draw.rect(screen, CARD_BG, (x, 20, 100, 50))
            screen.blit(FONT.render(txt, True, (0,0,0)), (x+10, 35))
            if selected_plant and txt.lower().startswith(selected_plant[:3]):
                pygame.draw.rect(screen, (255,255,0), (x, 20, 100, 50), 3)

        sun_txt = FONT.render(f"Sun: {sun_count}", True, UI_COLOR)
        screen.blit(sun_txt, (520, 35))

        for s in suns: s.draw(screen)
        for p in plants: p.draw(screen)
        for z in zombies: z.draw(screen)
        for b in bullets: b.draw(screen)

        if game_over:
            over_txt = BIG_FONT.render("ZOMBIES ATE YOUR BRAINS!", True, (200, 0, 0))
            rect = over_txt.get_rect(center=(WIDTH//2, HEIGHT//2))
            screen.blit(over_txt, rect)

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()