import pygame
import random

# 初始化
pygame.init()
SCREEN_WIDTH, SCREEN_HEIGHT = 800, 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("Mini Plants vs Zombies")
clock = pygame.time.Clock()

# 颜色定义
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
YELLOW = (255, 255, 0)

# --- 核心类定义 ---

class Plant(pygame.sprite.Sprite):
    def __init__(self, x, y, color):
        super().__init__()
        self.image = pygame.Surface((40, 40))
        self.image.fill(color)
        self.rect = self.image.get_rect(topleft=(x, y))
        self.health = 100

    def update(self):
        if self.health <= 0:
            self.kill()

class Peashooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, GREEN)
        self.shoot_timer = 0
        self.shoot_interval = 60  # 帧数

    def update(self, zombies, bullets):
        super().update()
        # 简单的同行检测：只要同行有僵尸就射击
        can_shoot = any(z.rect.y // 50 == self.rect.y // 50 and z.rect.x > self.rect.x for z in zombies)
        
        if can_shoot:
            self.shoot_timer += 1
            if self.shoot_timer >= self.shoot_interval:
                bullets.add(Bullet(self.rect.right, self.rect.centery))
                self.shoot_timer = 0

class Zombie(pygame.sprite.Sprite):
    def __init__(self, lane):
        super().__init__()
        self.image = pygame.Surface((30, 50))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.rect.x = SCREEN_WIDTH
        self.rect.y = lane * 50 + 10  # 假设每行高50
        self.health = 200
        self.speed = 1

    def update(self, plants):
        if self.health <= 0:
            self.kill()
            return

        # 碰撞检测：遇到植物就停止移动并攻击
        hit_plant = pygame.sprite.spritecollideany(self, plants)
        if hit_plant:
            # 简单攻击逻辑：每60帧扣10血
            if not hasattr(self, 'attack_timer'): self.attack_timer = 0
            self.attack_timer += 1
            if self.attack_timer >= 60:
                hit_plant.health -= 10
                self.attack_timer = 0
        else:
            self.rect.x -= self.speed

class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((10, 10))
        self.image.fill(YELLOW)
        self.rect = self.image.get_rect(center=(x, y))
        self.speed = 7

    def update(self):
        self.rect.x += self.speed
        if self.rect.x > SCREEN_WIDTH:
            self.kill()

# --- 游戏主循环 ---

plants = pygame.sprite.Group()
zombies = pygame.sprite.Group()
bullets = pygame.sprite.Group()

# 初始放置一个豌豆射手
plants.add(Peashooter(100, 110))

running = True
spawn_timer = 0

while running:
    screen.fill(WHITE)
    
    # 1. 事件处理
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            # 点击屏幕任意位置种植（简易演示）
            mx, my = pygame.mouse.get_pos()
            plants.add(Peashooter(mx - 20, my - 20)) 

    # 2. 逻辑更新
    spawn_timer += 1
    if spawn_timer > 180:  # 每3秒生成一个僵尸
        zombies.add(Zombie(random.randint(0, 10)))
        spawn_timer = 0

    plants.update(zombies, bullets)
    zombies.update(plants)
    bullets.update()

    # 子弹打僵尸
    for bullet in bullets:
        hits = pygame.sprite.spritecollide(bullet, zombies, False)
        if hits:
            for z in hits: z.health -= 20
            bullet.kill()

    # 3. 绘制
    plants.draw(screen)
    zombies.draw(screen)
    bullets.draw(screen)
    
    pygame.display.flip()
    clock.tick(60)

pygame.quit()