import pygame
import random
import time
import sys
import math
from enum import Enum

# 初始化pygame
pygame.init()

# 游戏常量
SCREEN_WIDTH = 1200
SCREEN_HEIGHT = 700
FPS = 60
GRID_SIZE = 80
GRID_WIDTH = 9
GRID_HEIGHT = 5
LAWN_X = 300
LAWN_Y = 100

# 颜色定义
GREEN = (50, 150, 50)
DARK_GREEN = (30, 100, 30)
LAWN_GREEN = (100, 200, 100)
SKY_BLUE = (135, 206, 235)
BROWN = (139, 69, 19)
YELLOW = (255, 215, 0)
RED = (255, 50, 50)
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
PURPLE = (128, 0, 128)
ORANGE = (255, 165, 0)
GRAY = (128, 128, 128)
LIGHT_GREEN = (144, 238, 144)

# 游戏状态
class GameState(Enum):
    MENU = 1
    PLAYING = 2
    PAUSED = 3
    GAME_OVER = 4

class PlantType(Enum):
    SUNFLOWER = 1
    PEASHOOTER = 2
    WALLNUT = 3
    CHERRY_BOMB = 4
    SNOW_PEA = 5
    REPEATER = 6

class ZombieType(Enum):
    NORMAL = 1
    CONEHEAD = 2
    BUCKETHEAD = 3
    NEWSPAPER = 4
    POLE_VAULTING = 5

class Sun:
    """太阳/金币"""
    def __init__(self, x, y, is_auto_collect=True):
        self.x = x
        self.y = y
        self.value = 25
        self.lifetime = 600  # 10秒
        self.is_auto_collect = is_auto_collect
        self.collect_timer = 0
        self.target_x = 20
        self.target_y = 20
        self.collecting = False
        self.speed = 10
        
    def update(self):
        if self.collecting:
            # 向太阳计数位置移动
            dx = self.target_x - self.x
            dy = self.target_y - self.y
            dist = math.sqrt(dx*dx + dy*dy)
            
            if dist < self.speed:
                return "collected"
            else:
                self.x += dx / dist * self.speed
                self.y += dy / dist * self.speed
        else:
            self.lifetime -= 1
            if self.lifetime <= 0:
                return "expired"
            
            if self.is_auto_collect:
                self.collect_timer += 1
                if self.collect_timer >= 300:  # 5秒后自动收集
                    self.collecting = True
        return "alive"
    
    def draw(self, screen):
        # 太阳主体
        pygame.draw.circle(screen, YELLOW, (int(self.x), int(self.y)), 15)
        
        # 太阳光芒
        for i in range(8):
            angle = math.radians(i * 45)
            x1 = self.x + math.cos(angle) * 20
            y1 = self.y + math.sin(angle) * 20
            x2 = self.x + math.cos(angle) * 30
            y2 = self.y + math.sin(angle) * 30
            pygame.draw.line(screen, YELLOW, (x1, y1), (x2, y2), 3)
        
        # 太阳面部
        pygame.draw.circle(screen, BLACK, (int(self.x-5), int(self.y-5)), 2)
        pygame.draw.circle(screen, BLACK, (int(self.x+5), int(self.y-5)), 2)
        pygame.draw.arc(screen, BLACK, (self.x-10, self.y, 20, 10), 0, math.pi, 2)
    
    def get_rect(self):
        return pygame.Rect(self.x - 20, self.y - 20, 40, 40)
    
    def start_collect(self):
        self.collecting = True

class Plant:
    """植物基类"""
    def __init__(self, x, y, plant_type):
        self.x = x
        self.y = y
        self.plant_type = plant_type
        self.health = 100
        self.cost = 100
        self.recharge_time = 750
        self.attack_timer = 0
        self.attack_speed = 100
        self.damage = 20
        self.range = SCREEN_WIDTH
        
        # 根据植物类型设置属性
        if plant_type == PlantType.SUNFLOWER:
            self.health = 50
            self.cost = 50
            self.recharge_time = 500
            self.sun_production_timer = 0
            self.sun_production_rate = 425
        elif plant_type == PlantType.PEASHOOTER:
            self.health = 100
            self.cost = 100
            self.recharge_time = 750
            self.attack_speed = 150
        elif plant_type == PlantType.WALLNUT:
            self.health = 400
            self.cost = 50
            self.recharge_time = 3000
        elif plant_type == PlantType.CHERRY_BOMB:
            self.health = 100
            self.cost = 150
            self.recharge_time = 5000
            self.explode_timer = 180
            self.exploding = False
        elif plant_type == PlantType.SNOW_PEA:
            self.health = 100
            self.cost = 175
            self.recharge_time = 1000
            self.attack_speed = 150
            self.slows = True
        elif plant_type == PlantType.REPEATER:
            self.health = 100
            self.cost = 200
            self.recharge_time = 1000
            self.attack_speed = 150
            self.double_shot = True
        
        self.rect = pygame.Rect(x, y, GRID_SIZE, GRID_SIZE)
    
    def draw(self, screen):
        if self.plant_type == PlantType.SUNFLOWER:
            # 向日葵
            pygame.draw.circle(screen, YELLOW, (self.x + GRID_SIZE//2, self.y + GRID_SIZE//2), 25)
            pygame.draw.circle(screen, (200, 150, 0), (self.x + GRID_SIZE//2, self.y + GRID_SIZE//2), 20)
            # 脸部
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2-10, self.y + GRID_SIZE//2-5), 3)
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2+10, self.y + GRID_SIZE//2-5), 3)
            pygame.draw.arc(screen, BLACK, (self.x + GRID_SIZE//2-15, self.y + GRID_SIZE//2+5, 30, 20), 0, math.pi, 2)
            
        elif self.plant_type == PlantType.PEASHOOTER:
            # 豌豆射手
            # 身体
            pygame.draw.rect(screen, GREEN, (self.x + 20, self.y + 20, 40, 40))
            pygame.draw.rect(screen, DARK_GREEN, (self.x + 20, self.y + 20, 40, 40), 2)
            # 头部
            pygame.draw.circle(screen, DARK_GREEN, (self.x + 40, self.y + 20), 15)
            # 嘴巴
            pygame.draw.rect(screen, BLACK, (self.x + 40, self.y + 15, 20, 10))
            # 叶子
            pygame.draw.polygon(screen, LIGHT_GREEN, [
                (self.x + 20, self.y + 30),
                (self.x + 10, self.y + 20),
                (self.x + 20, self.y + 20)
            ])
            pygame.draw.polygon(screen, LIGHT_GREEN, [
                (self.x + 60, self.y + 30),
                (self.x + 70, self.y + 20),
                (self.x + 60, self.y + 20)
            ])
            
        elif self.plant_type == PlantType.WALLNUT:
            # 坚果墙
            pygame.draw.circle(screen, BROWN, (self.x + GRID_SIZE//2, self.y + GRID_SIZE//2), 35)
            pygame.draw.circle(screen, (101, 67, 33), (self.x + GRID_SIZE//2, self.y + GRID_SIZE//2), 30)
            # 裂缝
            pygame.draw.line(screen, BLACK, (self.x + GRID_SIZE//2-15, self.y + GRID_SIZE//2-10), 
                           (self.x + GRID_SIZE//2+15, self.y + GRID_SIZE//2+10), 3)
            # 脸部
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2-10, self.y + GRID_SIZE//2-5), 4)
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2+10, self.y + GRID_SIZE//2-5), 4)
            pygame.draw.arc(screen, BLACK, (self.x + GRID_SIZE//2-15, self.y + GRID_SIZE//2+5, 30, 20), 0, math.pi, 3)
            
        elif self.plant_type == PlantType.CHERRY_BOMB:
            # 樱桃炸弹
            pygame.draw.circle(screen, RED, (self.x + GRID_SIZE//2-10, self.y + GRID_SIZE//2), 20)
            pygame.draw.circle(screen, RED, (self.x + GRID_SIZE//2+10, self.y + GRID_SIZE//2), 20)
            # 茎
            pygame.draw.line(screen, GREEN, (self.x + GRID_SIZE//2, self.y + GRID_SIZE//2-20), 
                           (self.x + GRID_SIZE//2, self.y + 10), 3)
            # 引线
            if self.exploding:
                pygame.draw.line(screen, ORANGE, (self.x + GRID_SIZE//2, self.y + 10), 
                               (self.x + GRID_SIZE//2, self.y + 5), 3)
            # 脸部
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2-15, self.y + GRID_SIZE//2-5), 3)
            pygame.draw.circle(screen, BLACK, (self.x + GRID_SIZE//2+5, self.y + GRID_SIZE//2-5), 3)
            pygame.draw.arc(screen, BLACK, (self.x + GRID_SIZE//2-20, self.y + GRID_SIZE//2+5, 40, 20), 0, math.pi, 2)
            
        elif self.plant_type == PlantType.SNOW_PEA:
            # 寒冰射手
            # 身体
            pygame.draw.rect(screen, (100, 150, 200), (self.x + 20, self.y + 20, 40, 40))
            pygame.draw.rect(screen, (50, 100, 150), (self.x + 20, self.y + 20, 40, 40), 2)
            # 头部
            pygame.draw.circle(screen, (50, 100, 150), (self.x + 40, self.y + 20), 15)
            # 嘴巴
            pygame.draw.rect(screen, (200, 230, 255), (self.x + 40, self.y + 15, 20, 10))
            # 冰晶
            pygame.draw.polygon(screen, (200, 230, 255), [
                (self.x + GRID_SIZE//2, self.y + 10),
                (self.x + GRID_SIZE//2-5, self.y + 20),
                (self.x + GRID_SIZE//2+5, self.y + 20)
            ])
            
        elif self.plant_type == PlantType.REPEATER:
            # 双重射手
            # 身体
            pygame.draw.rect(screen, (0, 150, 0), (self.x + 20, self.y + 20, 40, 40))
            pygame.draw.rect(screen, DARK_GREEN, (self.x + 20, self.y + 20, 40, 40), 2)
            # 双头
            pygame.draw.circle(screen, DARK_GREEN, (self.x + 30, self.y + 20), 12)
            pygame.draw.circle(screen, DARK_GREEN, (self.x + 50, self.y + 20), 12)
            # 嘴巴
            pygame.draw.rect(screen, BLACK, (self.x + 30, self.y + 15, 10, 8))
            pygame.draw.rect(screen, BLACK, (self.x + 50, self.y + 15, 10, 8))
            
    def update(self, current_time, zombies, projectiles, suns, grid_x, grid_y):
        if self.plant_type == PlantType.SUNFLOWER:
            # 向日葵产生阳光
            self.sun_production_timer += 1
            if self.sun_production_timer >= self.sun_production_rate:
                self.sun_production_timer = 0
                # 在向日葵旁边生成阳光
                sun_x = self.x + random.randint(-20, 20)
                sun_y = self.y + random.randint(-20, 20)
                suns.append(Sun(sun_x, sun_y, True))
                
        elif self.plant_type in [PlantType.PEASHOOTER, PlantType.SNOW_PEA, PlantType.REPEATER]:
            # 攻击型植物
            self.attack_timer += 1
            if self.attack_timer >= self.attack_speed:
                # 寻找目标
                target_zombie = None
                for zombie in zombies:
                    if (zombie.y == grid_y and  # 同一行
                        zombie.x > self.x and  # 在右侧
                        zombie.x - self.x < self.range):  # 在射程内
                        target_zombie = zombie
                        break
                
                if target_zombie:
                    self.attack_timer = 0
                    
                    if self.plant_type == PlantType.SNOW_PEA:
                        projectile_type = "snow"
                    else:
                        projectile_type = "pea"
                    
                    # 发射子弹
                    projectiles.append(Projectile(
                        self.x + GRID_SIZE,
                        self.y + GRID_SIZE//2,
                        projectile_type,
                        self.damage
                    ))
                    
                    if self.plant_type == PlantType.REPEATER:
                        # 双重射手发射两颗子弹
                        projectiles.append(Projectile(
                            self.x + GRID_SIZE + 10,
                            self.y + GRID_SIZE//2 + 5,
                            projectile_type,
                            self.damage
                        ))
                        
        elif self.plant_type == PlantType.CHERRY_BOMB:
            # 樱桃炸弹倒计时
            if self.exploding:
                self.explode_timer -= 1
                if self.explode_timer <= 0:
                    # 爆炸效果
                    explosion_radius = 100
                    for zombie in zombies[:]:
                        dist = math.sqrt((zombie.x - self.x)**2 + (zombie.y - self.y)**2)
                        if dist < explosion_radius:
                            zombie.health = 0
                    return "explode"
        return "alive"
    
    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0
    
    def get_rect(self):
        return pygame.Rect(self.x, self.y, GRID_SIZE, GRID_SIZE)

class Zombie:
    """僵尸基类"""
    def __init__(self, y, zombie_type=ZombieType.NORMAL):
        self.zombie_type = zombie_type
        self.x = SCREEN_WIDTH
        self.y = y
        self.speed = 0.5
        self.health = 100
        self.damage = 10
        self.attack_speed = 60
        self.attack_timer = 0
        self.eating = False
        self.frame = 0
        self.animation_timer = 0
        
        # 根据僵尸类型设置属性
        if zombie_type == ZombieType.NORMAL:
            self.health = 100
            self.speed = 0.5
        elif zombie_type == ZombieType.CONEHEAD:
            self.health = 250
            self.speed = 0.45
        elif zombie_type == ZombieType.BUCKETHEAD:
            self.health = 400
            self.speed = 0.4
        elif zombie_type == ZombieType.NEWSPAPER:
            self.health = 150
            self.speed = 0.6
            self.newspaper_health = 150
        elif zombie_type == ZombieType.POLE_VAULTING:
            self.health = 100
            self.speed = 1.0
            self.has_pole = True
            self.vaulting = False
        
        self.rect = pygame.Rect(self.x, self.y, GRID_SIZE, GRID_SIZE)
    
    def draw(self, screen):
        if self.zombie_type == ZombieType.NORMAL:
            # 普通僵尸
            # 身体
            pygame.draw.rect(screen, (100, 150, 100), (self.x + 20, self.y + 20, 40, 40))
            # 头
            pygame.draw.circle(screen, (200, 150, 100), (self.x + 40, self.y + 20), 15)
            # 脸
            pygame.draw.circle(screen, BLACK, (self.x + 35, self.y + 15), 3)
            pygame.draw.circle(screen, BLACK, (self.x + 45, self.y + 15), 3)
            pygame.draw.rect(screen, BLACK, (self.x + 35, self.y + 25, 10, 5))
            # 衣服
            pygame.draw.rect(screen, (150, 100, 100), (self.x + 20, self.y + 40, 40, 20))
            
        elif self.zombie_type == ZombieType.CONEHEAD:
            # 路障僵尸
            # 身体
            pygame.draw.rect(screen, (100, 150, 100), (self.x + 20, self.y + 20, 40, 40))
            # 头
            pygame.draw.circle(screen, (200, 150, 100), (self.x + 40, self.y + 20), 15)
            # 路障
            pygame.draw.rect(screen, ORANGE, (self.x + 20, self.y + 5, 40, 20))
            pygame.draw.rect(screen, (200, 100, 0), (self.x + 20, self.y + 5, 40, 20), 2)
            # 脸
            pygame.draw.circle(screen, BLACK, (self.x + 35, self.y + 15), 3)
            pygame.draw.circle(screen, BLACK, (self.x + 45, self.y + 15), 3)
            pygame.draw.rect(screen, BLACK, (self.x + 35, self.y + 25, 10, 5))
            
        elif self.zombie_type == ZombieType.BUCKETHEAD:
            # 铁桶僵尸
            # 身体
            pygame.draw.rect(screen, (100, 150, 100), (self.x + 20, self.y + 20, 40, 40))
            # 头
            pygame.draw.circle(screen, (200, 150, 100), (self.x + 40, self.y + 20), 15)
            # 铁桶
            pygame.draw.rect(screen, GRAY, (self.x + 20, self.y + 5, 40, 20))
            pygame.draw.rect(screen, (100, 100, 100), (self.x + 20, self.y + 5, 40, 20), 2)
            # 脸
            pygame.draw.circle(screen, BLACK, (self.x + 35, self.y + 15), 3)
            pygame.draw.circle(screen, BLACK, (self.x + 45, self.y + 15), 3)
            pygame.draw.rect(screen, BLACK, (self.x + 35, self.y + 25, 10, 5))
            
        elif self.zombie_type == ZombieType.NEWSPAPER:
            # 读报僵尸
            # 身体
            pygame.draw.rect(screen, (100, 150, 100), (self.x + 20, self.y + 20, 40, 40))
            # 头
            pygame.draw.circle(screen, (200, 150, 100), (self.x + 40, self.y + 20), 15)
            # 报纸
            pygame.draw.rect(screen, WHITE, (self.x + 10, self.y + 10, 60, 30))
            pygame.draw.rect(screen, BLACK, (self.x + 10, self.y + 10, 60, 30), 2)
            # 文字
            pygame.draw.rect(screen, BLACK, (self.x + 20, self.y + 15, 40, 2))
            pygame.draw.rect(screen, BLACK, (self.x + 20, self.y + 20, 40, 2))
            pygame.draw.rect(screen, BLACK, (self.x + 20, self.y + 25, 40, 2))
            
        elif self.zombie_type == ZombieType.POLE_VAULTING:
            # 撑杆僵尸
            # 身体
            pygame.draw.rect(screen, (100, 150, 100), (self.x + 20, self.y + 20, 40, 40))
            # 头
            pygame.draw.circle(screen, (200, 150, 100), (self.x + 40, self.y + 20), 15)
            # 杆子
            if self.has_pole:
                pygame.draw.rect(screen, BROWN, (self.x + 60, self.y + 20, 5, 40))
            # 脸
            pygame.draw.circle(screen, BLACK, (self.x + 35, self.y + 15), 3)
            pygame.draw.circle(screen, BLACK, (self.x + 45, self.y + 15), 3)
            pygame.draw.rect(screen, BLACK, (self.x + 35, self.y + 25, 10, 5))
        
        # 绘制生命条
        if self.health < (400 if self.zombie_type == ZombieType.BUCKETHEAD else 
                         250 if self.zombie_type == ZombieType.CONEHEAD else
                         150 if self.zombie_type == ZombieType.NEWSPAPER else 100):
            health_width = 40
            health_height = 5
            health_x = self.x + (GRID_SIZE - health_width) // 2
            health_y = self.y - 10
            
            # 计算生命百分比
            max_health = 400 if self.zombie_type == ZombieType.BUCKETHEAD else 250 if self.zombie_type == ZombieType.CONEHEAD else 150 if self.zombie_type == ZombieType.NEWSPAPER else 100
            health_percent = self.health / max_health
            
            # 背景
            pygame.draw.rect(screen, (100, 0, 0), (health_x, health_y, health_width, health_height))
            # 生命条
            pygame.draw.rect(screen, GREEN, (health_x, health_y, int(health_width * health_percent), health_height))
    
    def update(self, plants, grid_width, grid_height):
        if self.health <= 0:
            return "dead"
        
        # 动画
        self.animation_timer += 1
        if self.animation_timer >= 20:
            self.animation_timer = 0
            self.frame = (self.frame + 1) % 2
        
        # 检查是否在吃植物
        self.eating = False
        for plant in plants:
            plant_rect = plant.get_rect()
            zombie_rect = self.get_rect()
            
            if (abs(plant.y - self.y) < GRID_SIZE and
                abs(plant.x - self.x) < GRID_SIZE):
                # 在攻击范围内
                self.eating = True
                self.attack_timer += 1
                
                if self.attack_timer >= self.attack_speed:
                    self.attack_timer = 0
                    if plant.take_damage(self.damage):
                        return "plant_eaten"
                break
        
        if not self.eating:
            # 移动
            if self.zombie_type == ZombieType.POLE_VAULTING and self.has_pole:
                # 撑杆跳
                for plant in plants:
                    if (abs(plant.y - self.y) < GRID_SIZE and
                        plant.x - self.x < GRID_SIZE * 2 and
                        plant.x > self.x):
                        self.has_pole = False
                        self.x = plant.x + GRID_SIZE
                        self.vaulting = True
                        break
                
                if self.vaulting:
                    self.vaulting = False
                else:
                    self.x -= self.speed * 1.5
            else:
                self.x -= self.speed
            
            # 检查是否到达房子
            if self.x < LAWN_X - GRID_SIZE:
                return "reached_house"
        
        return "alive"
    
    def take_damage(self, damage, is_snow=False):
        if self.zombie_type == ZombieType.NEWSPAPER:
            if self.newspaper_health > 0:
                self.newspaper_health -= damage
                if self.newspaper_health <= 0:
                    self.speed = 1.0
            else:
                self.health -= damage
        else:
            self.health -= damage
            
        if is_snow:
            self.speed = max(0.1, self.speed * 0.5)
        
        return self.health <= 0
    
    def get_rect(self):
        return pygame.Rect(self.x, self.y, GRID_SIZE, GRID_SIZE)

class Projectile:
    """子弹/豌豆"""
    def __init__(self, x, y, proj_type, damage):
        self.x = x
        self.y = y
        self.type = proj_type
        self.speed = 5
        self.damage = damage
        
    def update(self):
        self.x += self.speed
        return self.x > SCREEN_WIDTH
    
    def draw(self, screen):
        if self.type == "pea":
            pygame.draw.circle(screen, GREEN, (int(self.x), int(self.y)), 5)
            pygame.draw.circle(screen, DARK_GREEN, (int(self.x), int(self.y)), 5, 1)
        else:  # snow
            pygame.draw.circle(screen, (200, 230, 255), (int(self.x), int(self.y)), 6)
            pygame.draw.circle(screen, (150, 200, 255), (int(self.x), int(self.y)), 6, 1)
    
    def get_rect(self):
        return pygame.Rect(self.x - 5, self.y - 5, 10, 10)

class PlantCard:
    """植物卡片"""
    def __init__(self, x, y, plant_type, cost):
        self.x = x
        self.y = y
        self.plant_type = plant_type
        self.cost = cost
        self.recharge_time = 0
        self.recharge_max = 750
        self.selected = False
        self.width = 60
        self.height = 80
        
    def draw(self, screen):
        # 卡片背景
        if self.selected:
            pygame.draw.rect(screen, (150, 200, 150), (self.x, self.y, self.width, self.height))
        else:
            pygame.draw.rect(screen, (200, 200, 200), (self.x, self.y, self.width, self.height))
        
        pygame.draw.rect(screen, (100, 100, 100), (self.x, self.y, self.width, self.height), 3)
        
        # 植物图标
        if self.plant_type == PlantType.SUNFLOWER:
            pygame.draw.circle(screen, YELLOW, (self.x + 30, self.y + 30), 15)
        elif self.plant_type == PlantType.PEASHOOTER:
            pygame.draw.rect(screen, GREEN, (self.x + 20, self.y + 20, 20, 20))
        elif self.plant_type == PlantType.WALLNUT:
            pygame.draw.circle(screen, BROWN, (self.x + 30, self.y + 30), 15)
        elif self.plant_type == PlantType.CHERRY_BOMB:
            pygame.draw.circle(screen, RED, (self.x + 30, self.y + 30), 15)
        elif self.plant_type == PlantType.SNOW_PEA:
            pygame.draw.rect(screen, (100, 150, 200), (self.x + 20, self.y + 20, 20, 20))
        elif self.plant_type == PlantType.REPEATER:
            pygame.draw.rect(screen, (0, 150, 0), (self.x + 20, self.y + 20, 20, 20))
        
        # 价格标签
        font = pygame.font.Font(None, 20)
        cost_text = font.render(str(self.cost), True, YELLOW)
        screen.blit(cost_text, (self.x + 30 - cost_text.get_width()//2, self.y + 55))
        
        # 冷却遮罩
        if self.recharge_time > 0:
            recharge_height = int(self.height * (self.recharge_time / self.recharge_max))
            pygame.draw.rect(screen, (0, 0, 0, 150), 
                           (self.x, self.y + self.height - recharge_height, 
                            self.width, recharge_height))
    
    def update(self):
        if self.recharge_time > 0:
            self.recharge_time -= 1
    
    def can_select(self, suns):
        return self.recharge_time <= 0 and suns >= self.cost
    
    def get_rect(self):
        return pygame.Rect(self.x, self.y, self.width, self.height)
    
    def select(self, suns):
        if self.can_select(suns):
            self.selected = True
            return self.cost
        return 0
    
    def unselect(self):
        self.selected = False

class PlantsVsZombiesGame:
    """植物大战僵尸游戏主类"""
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("植物大战僵尸 - 经典复刻版 - 卢思成制作")
        
        # 字体
        self.big_font = pygame.font.Font(None, 48)
        self.medium_font = pygame.font.Font(None, 36)
        self.small_font = pygame.font.Font(None, 24)
        
        # 游戏状态
        self.game_state = GameState.MENU
        self.clock = pygame.time.Clock()
        self.running = True
        
        # 游戏对象
        self.plants = []
        self.zombies = []
        self.projectiles = []
        self.suns = []
        self.plant_cards = []
        self.selected_plant = None
        self.hover_grid = None
        
        # 游戏参数
        self.suns_count = 50
        self.wave = 1
        self.zombies_spawned = 0
        self.zombies_killed = 0
        self.game_time = 0
        self.game_start_time = 0
        self.spawn_timer = 600  # 10秒后开始刷僵尸
        self.spawn_delay = 300
        self.zombies_per_wave = 5
        self.zombies_in_wave = 0
        
        # 开发者信息
        self.developer_info = {
            "name": "卢思成",
            "school": "东台市第一小学",
            "class": "六（10）班",
            "number": "19号"
        }
        
        # 初始化植物卡片
        self.init_plant_cards()
        
    def init_plant_cards(self):
        """初始化植物卡片"""
        plant_types = [
            (PlantType.SUNFLOWER, 50),
            (PlantType.PEASHOOTER, 100),
            (PlantType.WALLNUT, 50),
            (PlantType.SNOW_PEA, 175),
            (PlantType.REPEATER, 200),
            (PlantType.CHERRY_BOMB, 150)
        ]
        
        for i, (plant_type, cost) in enumerate(plant_types):
            x = 20
            y = 100 + i * 90
            self.plant_cards.append(PlantCard(x, y, plant_type, cost))
    
    def get_grid_position(self, x, y):
        """获取网格位置"""
        if (LAWN_X <= x < LAWN_X + GRID_WIDTH * GRID_SIZE and
            LAWN_Y <= y < LAWN_Y + GRID_HEIGHT * GRID_SIZE):
            grid_x = (x - LAWN_X) // GRID_SIZE
            grid_y = (y - LAWN_Y) // GRID_SIZE
            return (grid_x, grid_y)
        return None
    
    def spawn_zombie(self):
        """生成僵尸"""
        if self.zombies_in_wave >= self.zombies_per_wave:
            return
        
        # 随机选择行
        row = random.randint(0, GRID_HEIGHT - 1)
        y = LAWN_Y + row * GRID_SIZE
        
        # 根据波数选择僵尸类型
        if self.wave < 3:
            zombie_types = [ZombieType.NORMAL]
        elif self.wave < 5:
            zombie_types = [ZombieType.NORMAL, ZombieType.CONEHEAD]
        elif self.wave < 7:
            zombie_types = [ZombieType.NORMAL, ZombieType.CONEHEAD, ZombieType.BUCKETHEAD]
        else:
            zombie_types = [ZombieType.NORMAL, ZombieType.CONEHEAD, 
                          ZombieType.BUCKETHEAD, ZombieType.NEWSPAPER, ZombieType.POLE_VAULTING]
        
        zombie_type = random.choice(zombie_types)
        self.zombies.append(Zombie(y, zombie_type))
        
        self.zombies_spawned += 1
        self.zombies_in_wave += 1
    
    def spawn_sun(self):
        """生成阳光"""
        if random.random() < 0.01:  # 1%几率每帧生成
            x = random.randint(LAWN_X, SCREEN_WIDTH - 100)
            y = random.randint(LAWN_Y, SCREEN_HEIGHT - 100)
            self.suns.append(Sun(x, y, True))
    
    def update(self):
        """更新游戏逻辑"""
        if self.game_state != GameState.PLAYING:
            return
        
        self.game_time += 1
        
        # 更新植物卡片
        for card in self.plant_cards:
            card.update()
        
        # 自动生成阳光
        self.spawn_sun()
        
        # 更新阳光
        for sun in self.suns[:]:
            status = sun.update()
            if status == "collected":
                self.suns_count += sun.value
                self.suns.remove(sun)
            elif status == "expired":
                self.suns.remove(sun)
        
        # 更新植物
        for plant in self.plants[:]:
            grid_x = (plant.x - LAWN_X) // GRID_SIZE
            grid_y = (plant.y - LAWN_Y) // GRID_SIZE
            status = plant.update(
                self.game_time,
                self.zombies,
                self.projectiles,
                self.suns,
                grid_x,
                grid_y
            )
            if status == "explode":
                self.plants.remove(plant)
        
        # 更新子弹
        for projectile in self.projectiles[:]:
            if projectile.update():
                self.projectiles.remove(projectile)
                continue
            
            # 检查子弹是否击中僵尸
            for zombie in self.zombies[:]:
                if projectile.get_rect().colliderect(zombie.get_rect()):
                    is_snow = projectile.type == "snow"
                    if zombie.take_damage(projectile.damage, is_snow):
                        if zombie.health <= 0:
                            self.zombies_killed += 1
                            self.suns_count += 25
                            self.zombies.remove(zombie)
                    if projectile in self.projectiles:
                        self.projectiles.remove(projectile)
                    break
        
        # 更新僵尸
        for zombie in self.zombies[:]:
            status = zombie.update(self.plants, GRID_WIDTH, GRID_HEIGHT)
            if status == "dead":
                self.zombies_killed += 1
                self.suns_count += 25
                self.zombies.remove(zombie)
            elif status == "reached_house":
                self.game_state = GameState.GAME_OVER
        
        # 生成僵尸
        if self.game_time > self.spawn_timer:
            if self.game_time % self.spawn_delay == 0:
                self.spawn_zombie()
            
            # 如果所有僵尸都被消灭，开始下一波
            if len(self.zombies) == 0 and self.zombies_in_wave >= self.zombies_per_wave:
                self.wave += 1
                self.zombies_in_wave = 0
                self.zombies_per_wave = min(20, self.zombies_per_wave + 2)
                self.spawn_delay = max(60, self.spawn_delay - 20)  # 加快生成速度
    
    def draw_background(self):
        """绘制游戏背景"""
        # 天空
        self.screen.fill(SKY_BLUE)
        
        # 草地
        for row in range(GRID_HEIGHT):
            for col in range(GRID_WIDTH):
                x = LAWN_X + col * GRID_SIZE
                y = LAWN_Y + row * GRID_SIZE
                pygame.draw.rect(self.screen, LAWN_GREEN, (x, y, GRID_SIZE, GRID_SIZE))
                pygame.draw.rect(self.screen, DARK_GREEN, (x, y, GRID_SIZE, GRID_SIZE), 1)
        
        # 房子
        pygame.draw.rect(self.screen, BROWN, (LAWN_X - 100, LAWN_Y, 100, GRID_HEIGHT * GRID_SIZE))
        pygame.draw.rect(self.screen, (101, 67, 33), (LAWN_X - 100, LAWN_Y, 100, GRID_HEIGHT * GRID_SIZE), 3)
        
        # 屋顶
        roof_points = [
            (LAWN_X - 100, LAWN_Y),
            (LAWN_X - 50, LAWN_Y - 50),
            (LAWN_X, LAWN_Y)
        ]
        pygame.draw.polygon(self.screen, (150, 100, 50), roof_points)
        
        # 门
        pygame.draw.rect(self.screen, (200, 150, 100), 
                        (LAWN_X - 70, LAWN_Y + GRID_HEIGHT * GRID_SIZE//2 - 30, 40, 60))
        pygame.draw.circle(self.screen, BLACK, (LAWN_X - 45, LAWN_Y + GRID_HEIGHT * GRID_SIZE//2), 5)
    
    def draw_ui(self):
        """绘制游戏UI"""
        # 阳光数量
        sun_text = self.medium_font.render(f"阳光: {self.suns_count}", True, YELLOW)
        self.screen.blit(sun_text, (20, 20))
        
        # 波数
        wave_text = self.medium_font.render(f"第 {self.wave} 波", True, WHITE)
        self.screen.blit(wave_text, (SCREEN_WIDTH - 150, 20))
        
        # 消灭僵尸数
        kills_text = self.small_font.render(f"消灭僵尸: {self.zombies_killed}", True, GREEN)
        self.screen.blit(kills_text, (SCREEN_WIDTH - 150, 60))
        
        # 游戏时间
        minutes = self.game_time // 3600
        seconds = (self.game_time // 60) % 60
        time_text = self.small_font.render(f"时间: {minutes:02d}:{seconds:02d}", True, WHITE)
        self.screen.blit(time_text, (SCREEN_WIDTH - 150, 90))
        
        # 植物卡片
        for card in self.plant_cards:
            card.draw(self.screen)
        
        # 网格高亮
        if self.hover_grid and self.selected_plant:
            grid_x, grid_y = self.hover_grid
            x = LAWN_X + grid_x * GRID_SIZE
            y = LAWN_Y + grid_y * GRID_SIZE
            
            # 检查是否可以放置
            can_place = True
            for plant in self.plants:
                if (plant.x == x and plant.y == y):
                    can_place = False
                    break
            
            if can_place:
                # 找到对应的卡片
                for card in self.plant_cards:
                    if card.plant_type == self.selected_plant:
                        if self.suns_count >= card.cost:
                            color = GREEN
                        else:
                            color = RED
                        break
            else:
                color = RED
            
            pygame.draw.rect(self.screen, color, (x, y, GRID_SIZE, GRID_SIZE), 3)
    
    def draw_menu(self):
        """绘制主菜单"""
        # 背景
        self.screen.fill(SKY_BLUE)
        
        # 标题
        title = self.big_font.render("植物大战僵尸", True, GREEN)
        title_shadow = self.big_font.render("植物大战僵尸", True, DARK_GREEN)
        self.screen.blit(title_shadow, (SCREEN_WIDTH//2 - title.get_width()//2 + 3, 103))
        self.screen.blit(title, (SCREEN_WIDTH//2 - title.get_width()//2, 100))
        
        # 副标题
        subtitle = self.medium_font.render("经典复刻版", True, YELLOW)
        self.screen.blit(subtitle, (SCREEN_WIDTH//2 - subtitle.get_width()//2, 160))
        
        # 开始按钮
        button_width, button_height = 200, 60
        button_x = SCREEN_WIDTH//2 - button_width//2
        button_y = 250
        
        pygame.draw.rect(self.screen, GREEN, 
                       (button_x, button_y, button_width, button_height))
        pygame.draw.rect(self.screen, DARK_GREEN, 
                       (button_x, button_y, button_width, button_height), 3)
        
        start_text = self.medium_font.render("开始游戏", True, WHITE)
        self.screen.blit(start_text, (button_x + button_width//2 - start_text.get_width()//2, 
                                    button_y + button_height//2 - start_text.get_height()//2))
        
        # 操作说明
        instructions = [
            "游戏说明:",
            "1. 选择植物卡片放置在草地上",
            "2. 收集阳光购买更多植物",
            "3. 阻止僵尸进入房子",
            "4. 10秒后开始出现僵尸",
            "5. 难度会逐渐增加",
            "",
            "特殊功能:",
            "✓ 阳光自动收集",
            "✓ 僵尸波数递增",
            "✓ 多种植物和僵尸类型"
        ]
        
        y_offset = 350
        for i, line in enumerate(instructions):
            if i == 0:
                text = self.medium_font.render(line, True, YELLOW)
            else:
                text = self.small_font.render(line, True, WHITE)
            self.screen.blit(text, (SCREEN_WIDTH//2 - text.get_width()//2, y_offset + i * 30))
        
        # 开发者信息
        dev_text = f"开发者: {self.developer_info['name']}  {self.developer_info['school']} {self.developer_info['class']} {self.developer_info['number']}"
        dev_render = self.small_font.render(dev_text, True, (150, 150, 150))
        self.screen.blit(dev_render, (SCREEN_WIDTH//2 - dev_render.get_width()//2, 
                                    SCREEN_HEIGHT - 50))
        
        return pygame.Rect(button_x, button_y, button_width, button_height)
    
    def draw_game_over(self):
        """绘制游戏结束界面"""
        # 半透明覆盖层
        overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT), pygame.SRCALPHA)
        overlay.fill((0, 0, 0, 150))
        self.screen.blit(overlay, (0, 0))
        
        # 游戏结束标题
        game_over = self.big_font.render("游戏结束", True, RED)
        self.screen.blit(game_over, (SCREEN_WIDTH//2 - game_over.get_width()//2, 150))
        
        # 统计信息
        minutes = self.game_time // 3600
        seconds = (self.game_time // 60) % 60
        
        stats = [
            f"坚持波数: {self.wave}",
            f"消灭僵尸: {self.zombies_killed}",
            f"游戏时间: {minutes:02d}:{seconds:02d}",
            f"获得阳光: {self.suns_count}"
        ]
        
        y_offset = 220
        for i, stat in enumerate(stats):
            stat_render = self.medium_font.render(stat, True, WHITE)
            self.screen.blit(stat_render, (SCREEN_WIDTH//2 - stat_render.get_width()//2, y_offset + i * 40))
        
        # 重新开始按钮
        button_width, button_height = 200, 50
        button_x = SCREEN_WIDTH//2 - button_width//2
        button_y = 400
        
        pygame.draw.rect(self.screen, GREEN, 
                        (button_x, button_y, button_width, button_height))
        pygame.draw.rect(self.screen, DARK_GREEN, 
                        (button_x, button_y, button_width, button_height), 3)
        
        restart_text = self.medium_font.render("重新开始", True, WHITE)
        self.screen.blit(restart_text, (button_x + button_width//2 - restart_text.get_width()//2, 
                                      button_y + button_height//2 - restart_text.get_height()//2))
        
        # 返回菜单按钮
        menu_button_y = button_y + button_height + 20
        pygame.draw.rect(self.screen, SKY_BLUE, 
                        (button_x, menu_button_y, button_width, button_height))
        pygame.draw.rect(self.screen, (100, 150, 200), 
                        (button_x, menu_button_y, button_width, button_height), 3)
        
        menu_text = self.medium_font.render("返回菜单", True, WHITE)
        self.screen.blit(menu_text, (button_x + button_width//2 - menu_text.get_width()//2, 
                                   menu_button_y + button_height//2 - menu_text.get_height()//2))
        
        return (pygame.Rect(button_x, button_y, button_width, button_height),
                pygame.Rect(button_x, menu_button_y, button_width, button_height))
    
    def start_game(self):
        """开始新游戏"""
        self.plants = []
        self.zombies = []
        self.projectiles = []
        self.suns = []
        
        self.suns_count = 50
        self.wave = 1
        self.zombies_spawned = 0
        self.zombies_killed = 0
        self.game_time = 0
        self.spawn_timer = 600
        self.spawn_delay = 300
        self.zombies_per_wave = 5
        self.zombies_in_wave = 0
        
        self.selected_plant = None
        for card in self.plant_cards:
            card.unselect()
        
        self.game_state = GameState.PLAYING
    
    def handle_events(self):
        """处理事件"""
        mouse_pos = pygame.mouse.get_pos()
        self.hover_grid = self.get_grid_position(mouse_pos[0], mouse_pos[1])
        
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                self.running = False
                pygame.quit()
                sys.exit()
            
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if self.game_state == GameState.MENU:
                    start_button = self.draw_menu()
                    if start_button.collidepoint(mouse_pos):
                        self.start_game()
                
                elif self.game_state == GameState.PLAYING:
                    # 检查是否点击植物卡片
                    for card in self.plant_cards:
                        if card.get_rect().collidepoint(mouse_pos):
                            if card.can_select(self.suns_count):
                                cost = card.select(self.suns_count)
                                if cost > 0:
                                    self.selected_plant = PlantType(card.plant_type.value)
                                    # 取消选择其他卡片
                                    for other_card in self.plant_cards:
                                        if other_card != card:
                                            other_card.unselect()
                            break
                    
                    # 检查是否点击网格放置植物
                    if self.selected_plant and self.hover_grid:
                        grid_x, grid_y = self.hover_grid
                        x = LAWN_X + grid_x * GRID_SIZE
                        y = LAWN_Y + grid_y * GRID_SIZE
                        
                        # 检查是否可以放置
                        can_place = True
                        for plant in self.plants:
                            if (plant.x == x and plant.y == y):
                                can_place = False
                                break
                        
                        if can_place:
                            # 找到对应的卡片
                            for card in self.plant_cards:
                                if card.plant_type == self.selected_plant:
                                    if self.suns_count >= card.cost:
                                        # 创建植物
                                        new_plant = Plant(x, y, self.selected_plant)
                                        self.plants.append(new_plant)
                                        self.suns_count -= card.cost
                                        card.recharge_time = card.recharge_max
                                        card.unselect()
                                        self.selected_plant = None
                                    break
                    
                    # 检查是否点击阳光
                    for sun in self.suns[:]:
                        if sun.get_rect().collidepoint(mouse_pos):
                            sun.start_collect()
                
                elif self.game_state == GameState.GAME_OVER:
                    restart_button, menu_button = self.draw_game_over()
                    if restart_button.collidepoint(mouse_pos):
                        self.start_game()
                    elif menu_button.collidepoint(mouse_pos):
                        self.game_state = GameState.MENU
            
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    if self.game_state == GameState.PLAYING:
                        self.game_state = GameState.MENU
                    elif self.game_state == GameState.MENU:
                        self.running = False
                elif event.key == pygame.K_SPACE and self.game_state == GameState.MENU:
                    self.start_game()
                elif event.key == pygame.K_r and self.game_state == GameState.GAME_OVER:
                    self.start_game()
    
    def draw(self):
        """绘制游戏"""
        if self.game_state == GameState.MENU:
            self.draw_menu()
            
        elif self.game_state == GameState.PLAYING:
            # 绘制背景
            self.draw_background()
            
            # 绘制植物
            for plant in self.plants:
                plant.draw(self.screen)
            
            # 绘制僵尸
            for zombie in self.zombies:
                zombie.draw(self.screen)
            
            # 绘制子弹
            for projectile in self.projectiles:
                projectile.draw(self.screen)
            
            # 绘制阳光
            for sun in self.suns:
                sun.draw(self.screen)
            
            # 绘制UI
            self.draw_ui()
            
        elif self.game_state == GameState.GAME_OVER:
            # 绘制游戏画面
            self.draw_background()
            
            for plant in self.plants:
                plant.draw(self.screen)
            
            for zombie in self.zombies:
                zombie.draw(self.screen)
            
            for projectile in self.projectiles:
                projectile.draw(self.screen)
            
            for sun in self.suns:
                sun.draw(self.screen)
            
            self.draw_ui()
            self.draw_game_over()
    
    def run(self):
        """运行游戏"""
        print("""
        植物大战僵尸 - 经典复刻版
        开发者: 卢思成
        学校: 东台市第一小学
        班级: 六（10）班
        学号: 19号
        
        游戏特色:
        1. 阳光自动收集
        2. 10秒后开始刷新僵尸
        3. 难度逐渐增加
        4. 多种植物和僵尸类型
        5. 完全复刻原版玩法
        
        操作说明:
        1. 选择植物卡片放置在草地上
        2. 点击阳光立即收集
        3. 阻止僵尸进入房子
        4. ESC键返回菜单
        """)
        
        while self.running:
            self.handle_events()
            self.update()
            self.draw()
            pygame.display.flip()
            self.clock.tick(FPS)
        
        pygame.quit()
        sys.exit()

# 运行游戏
if __name__ == "__main__":
    try:
        game = PlantsVsZombiesGame()
        game.run()
    except Exception as e:
        print(f"游戏启动出错: {e}")
        import traceback
        traceback.print_exc()
        input("按回车键退出...")