import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏窗口设置
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("植物大战僵尸 - 简化版")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
BROWN = (139, 69, 19)
YELLOW = (255, 255, 0)

# 游戏设置
FPS = 60
clock = pygame.time.Clock()

class Sunflower:
    """向日葵类 - 生产阳光"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 50
        self.height = 50
        self.production_timer = 0
        self.production_interval = 180  # 3秒生产一个阳光
        
    def draw(self, screen):
        # 绘制向日葵身体
        pygame.draw.circle(screen, YELLOW, (self.x + 25, self.y + 25), 20)
        pygame.draw.circle(screen, BROWN, (self.x + 25, self.y + 35), 10)
        # 绘制花瓣
        for i in range(8):
            angle = i * 45
            x_offset = int(20 * pygame.math.Vector2(1, 0).rotate(angle).x)
            y_offset = int(20 * pygame.math.Vector2(1, 0).rotate(angle).y)
            pygame.draw.circle(screen, YELLOW, (self.x + 25 + x_offset, self.y + 25 + y_offset), 8)
    
    def update(self):
        self.production_timer += 1
        if self.production_timer >= self.production_interval:
            self.production_timer = 0
            return Sun(self.x + 25, self.y + 25)
        return None

class PeaShooter:
    """豌豆射手类"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 50
        self.height = 50
        self.shoot_timer = 0
        self.shoot_interval = 60  # 1秒射击一次
        
    def draw(self, screen):
        # 绘制豌豆射手身体
        pygame.draw.circle(screen, GREEN, (self.x + 25, self.y + 25), 20)
        # 绘制嘴巴（枪口）
        pygame.draw.circle(screen, RED, (self.x + 45, self.y + 25), 5)
        # 绘制眼睛
        pygame.draw.circle(screen, BLACK, (self.x + 20, self.y + 20), 3)
        
    def update(self):
        self.shoot_timer += 1
        if self.shoot_timer >= self.shoot_interval:
            self.shoot_timer = 0
            return Pea(self.x + 50, self.y + 25)
        return None

class Pea:
    """豌豆子弹类"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 10
        self.height = 10
        self.speed = 5
        
    def draw(self, screen):
        pygame.draw.circle(screen, GREEN, (int(self.x), int(self.y)), 5)
        
    def update(self):
        self.x += self.speed
        return self.x < SCREEN_WIDTH

class Sun:
    """阳光类"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.width = 20
        self.height = 20
        self.lifetime = 300  # 5秒后消失
        
    def draw(self, screen):
        pygame.draw.circle(screen, YELLOW, (int(self.x), int(self.y)), 10)
        pygame.draw.circle(screen, (255, 200, 0), (int(self.x), int(self.y)), 8)
        
    def update(self):
        self.lifetime -= 1
        return self.lifetime > 0

class Zombie:
    """僵尸类"""
    def __init__(self, y):
        self.x = SCREEN_WIDTH
        self.y = y
        self.width = 50
        self.height = 50
        self.health = 5
        self.speed = 1
        self.attack_timer = 0
        
    def draw(self, screen):
        # 绘制僵尸
        pygame.draw.rect(screen, (100, 100, 100), (self.x, self.y, self.width, self.height))
        pygame.draw.circle(screen, BLACK, (self.x + 40, self.y + 15), 3)
        pygame.draw.circle(screen, BLACK, (self.x + 20, self.y + 15), 3)
        # 绘制血条
        bar_width = self.width
        bar_height = 5
        bar_x = self.x
        bar_y = self.y - 10
        health_percent = self.health / 5
        pygame.draw.rect(screen, RED, (bar_x, bar_y, bar_width, bar_height))
        pygame.draw.rect(screen, GREEN, (bar_x, bar_y, bar_width * health_percent, bar_height))
        
    def update(self):
        self.x -= self.speed
        
    def take_damage(self, damage):
        self.health -= damage
        return self.health <= 0

class Game:
    """游戏主类"""
    def __init__(self):
        self.plants = []
        self.zombies = []
        self.peas = []
        self.suns = []
        self.sun_points = 150
        self.selected_plant = None
        self.font = pygame.font.Font(None, 36)
        self.grid = [[None for _ in range(5)] for _ in range(9)]
        self.game_over = False
        
    def draw_grid(self, screen):
        """绘制网格线"""
        for x in range(9):
            for y in range(5):
                rect = pygame.Rect(80 + x * 80, 100 + y * 100, 80, 100)
                pygame.draw.rect(screen, (200, 200, 200), rect, 1)
                
    def draw_ui(self, screen):
        """绘制UI"""
        # 显示阳光点数
        sun_text = self.font.render(f"阳光: {self.sun_points}", True, YELLOW)
        screen.blit(sun_text, (10, 10))
        
        # 显示植物卡片
        plant_cards = [
            ("向日葵", 50, 80),
            ("豌豆射手", 50, 130)
        ]
        
        for i, (name, x, y) in enumerate(plant_cards):
            rect = pygame.Rect(x, y, 60, 40)
            if i == 0:
                color = YELLOW
            else:
                color = GREEN
            pygame.draw.rect(screen, color, rect)
            pygame.draw.rect(screen, BLACK, rect, 2)
            text = self.font.render(name[:2], True, BLACK)
            screen.blit(text, (x + 5, y + 10))
            
    def handle_click(self, pos):
        """处理点击事件"""
        x, y = pos
        
        # 检查是否点击了植物卡片
        if 50 <= x <= 110:
            if 80 <= y <= 120:  # 向日葵
                if self.sun_points >= 50:
                    self.selected_plant = "sunflower"
                    self.sun_points -= 50
            elif 130 <= y <= 170:  # 豌豆射手
                if self.sun_points >= 100:
                    self.selected_plant = "peashooter"
                    self.sun_points -= 100
            return
            
        # 检查是否点击了网格放置植物
        if self.selected_plant and x >= 80 and x <= 800 and y >= 100 and y <= 600:
            grid_x = (x - 80) // 80
            grid_y = (y - 100) // 100
            
            if 0 <= grid_x < 9 and 0 <= grid_y < 5:
                if self.grid[grid_x][grid_y] is None:
                    plant_x = 80 + grid_x * 80
                    plant_y = 100 + grid_y * 100
                    
                    if self.selected_plant == "sunflower":
                        self.plants.append(Sunflower(plant_x, plant_y))
                        self.grid[grid_x][grid_y] = "sunflower"
                    elif self.selected_plant == "peashooter":
                        self.plants.append(PeaShooter(plant_x, plant_y))
                        self.grid[grid_x][grid_y] = "peashooter"
                    
                    self.selected_plant = None
        
        # 检查是否点击了阳光
        for sun in self.suns[:]:
            if (sun.x - 10 <= x <= sun.x + 10 and 
                sun.y - 10 <= y <= sun.y + 10):
                self.suns.remove(sun)
                self.sun_points += 25
                
    def spawn_zombie(self):
        """生成僵尸"""
        if random.randint(1, 100) <= 2:  # 2%的概率生成僵尸
            y = random.randint(0, 4)
            self.zombies.append(Zombie(100 + y * 100))
            
    def update(self):
        """更新游戏状态"""
        if self.game_over:
            return
            
        # 更新植物
        for plant in self.plants[:]:
            if isinstance(plant, Sunflower):
                new_sun = plant.update()
                if new_sun:
                    self.suns.append(new_sun)
            elif isinstance(plant, PeaShooter):
                new_pea = plant.update()
                if new_pea:
                    self.peas.append(new_pea)
                    
        # 更新阳光
        for sun in self.suns[:]:
            if not sun.update():
                self.suns.remove(sun)
                
        # 更新豌豆
        for pea in self.peas[:]:
            if not pea.update():
                self.peas.remove(pea)
                
        # 更新僵尸
        for zombie in self.zombies[:]:
            zombie.update()
            
            # 检查僵尸是否到达房子
            if zombie.x < 80:
                self.game_over = True
                
        # 检查豌豆和僵尸的碰撞
        for pea in self.peas[:]:
            for zombie in self.zombies[:]:
                if (zombie.x < pea.x < zombie.x + zombie.width and
                    zombie.y < pea.y < zombie.y + zombie.height):
                    if zombie.take_damage(1):
                        self.zombies.remove(zombie)
                    self.peas.remove(pea)
                    break
                    
        # 生成新僵尸
        self.spawn_zombie()
        
    def draw(self, screen):
        """绘制所有元素"""
        screen.fill(WHITE)
        
        # 绘制网格
        self.draw_grid(screen)
        
        # 绘制UI
        self.draw_ui(screen)
        
        # 绘制植物
        for plant in self.plants:
            plant.draw(screen)
            
        # 绘制僵尸
        for zombie in self.zombies:
            zombie.draw(screen)
            
        # 绘制豌豆
        for pea in self.peas:
            pea.draw(screen)
            
        # 绘制阳光
        for sun in self.suns:
            sun.draw(screen)
            
        # 显示选中的植物
        if self.selected_plant:
            x, y = pygame.mouse.get_pos()
            pygame.draw.circle(screen, (255, 255, 0, 128), (x, y), 25, 2)
            
        # 游戏结束画面
        if self.game_over:
            overlay = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
            overlay.set_alpha(128)
            overlay.fill(BLACK)
            screen.blit(overlay, (0, 0))
            
            game_over_text = self.font.render("游戏结束! 按R重新开始", True, RED)
            text_rect = game_over_text.get_rect(center=(SCREEN_WIDTH//2, SCREEN_HEIGHT//2))
            screen.blit(game_over_text, text_rect)
            
        pygame.display.flip()
        
    def reset(self):
        """重置游戏"""
        self.plants.clear()
        self.zombies.clear()
        self.peas.clear()
        self.suns.clear()
        self.sun_points = 150
        self.selected_plant = None
        self.grid = [[None for _ in range(5)] for _ in range(9)]
        self.game_over = False
        
def main():
    """主函数"""
    game = Game()
    running = True
    
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 1:  # 左键点击
                    if game.game_over:
                        game.reset()
                    else:
                        game.handle_click(event.pos)
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r and game.game_over:
                    game.reset()
                    
        game.update()
        game.draw(screen)
        clock.tick(FPS)
        
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()