import pygame
import random
import sys

# ===================== 游戏常量 =====================
WIDTH, HEIGHT = 900, 600
FPS = 60
GRID_SIZE = 80  # 格子大小
ROW_COUNT = 5   # 5行草地
SUN_COST_SUNFLOWER = 50
SUN_COST_PEASHOOTER = 100

# 颜色定义
GRASS_COLOR = (80, 160, 60)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
YELLOW = (255, 220, 0)
GREEN_PLANT = (20, 180, 30)
ZOMBIE_COLOR = (100, 70, 40)
PEA_COLOR = (0, 200, 255)
RED = (255, 0, 0)

# 初始化Pygame
pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("简易植物大战僵尸 Python版")
clock = pygame.time.Clock()
# 修复字体：替换simhei为Arial，无中文环境也不会报错
try:
    font = pygame.font.SysFont("Arial", 30)
except:
    # 兜底默认字体，绝对不会报错
    font = pygame.font.Font(None, 30)

# ===================== 游戏对象类 =====================
class Sun:
    """阳光道具"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = 15
        self.life = 300  # 存活帧数

    def draw(self):
        pygame.draw.circle(screen, YELLOW, (int(self.x), int(self.y)), self.radius)

class SunFlower:
    """向日葵：产生阳光"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.w = GRID_SIZE
        self.h = GRID_SIZE
        self.hp = 100
        self.timer = 180  # 每3秒产阳光

    def update(self, sun_list):
        self.timer -= 1
        if self.timer <= 0:
            sun_list.append(Sun(self.x+40, self.y-20))
            self.timer = 180

    def draw(self):
        pygame.draw.rect(screen, YELLOW, (self.x, self.y, self.w, self.h))
        pygame.draw.circle(screen, BLACK, (self.x+40, self.y+40), 20)

class PeaShooter:
    """豌豆射手：发射子弹"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.w = GRID_SIZE
        self.h = GRID_SIZE
        self.hp = 100
        self.shoot_timer = 60  # 1秒一发豌豆

    def update(self, bullet_list, zombies):
        self.shoot_timer -= 1
        # 同行有僵尸才射击
        same_row_zombie = False
        for z in zombies:
            if abs(z.y - self.y) < 10:
                same_row_zombie = True
                break
        if self.shoot_timer <= 0 and same_row_zombie:
            bullet_list.append(Bullet(self.x+80, self.y+40))
            self.shoot_timer = 60

    def draw(self):
        pygame.draw.rect(screen, GREEN_PLANT, (self.x, self.y, self.w, self.h))

class Bullet:
    """豌豆子弹"""
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.r = 8
        self.speed = 6
        self.damage = 25

    def update(self):
        self.x += self.speed

    def draw(self):
        pygame.draw.circle(screen, PEA_COLOR, (int(self.x), int(self.y)), self.r)

class Zombie:
    """僵尸"""
    def __init__(self, row):
        self.y = row * GRID_SIZE
        self.x = WIDTH
        self.w = GRID_SIZE
        self.h = GRID_SIZE
        self.hp = 100
        self.speed = 0.3
        self.eat_timer = 60  # 啃植物冷却

    def update(self, plants):
        # 检测前方植物
        hit_plant = None
        for p in plants:
            if abs(p.y - self.y) < 10 and p.x < self.x:
                hit_plant = p
                break
        if hit_plant:
            self.eat_timer -= 1
            if self.eat_timer <= 0:
                hit_plant.hp -= 10
                self.eat_timer = 60
        else:
            self.x -= self.speed

    def draw(self):
        pygame.draw.rect(screen, ZOMBIE_COLOR, (self.x, self.y, self.w, self.h))

# ===================== 游戏主逻辑 =====================
def main():
    sun_count = 50  # 初始阳光
    plants = []     # 所有植物
    bullets = []    # 子弹列表
    zombies = []    # 僵尸列表
    sun_items = []  # 掉落阳光
    spawn_zombie_timer = 300  # 僵尸生成间隔
    game_over = False
    select_plant = 0  # 0向日葵 / 1豌豆射手

    while True:
        screen.fill(GRASS_COLOR)
        # 绘制草地网格
        for r in range(ROW_COUNT):
            pygame.draw.line(screen, BLACK, (0, r*GRID_SIZE), (WIDTH, r*GRID_SIZE), 2)

        # 事件监听
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()
            if game_over:
                continue
            # 鼠标点击收集阳光 / 种植植物
            if event.type == pygame.MOUSEBUTTONDOWN:
                mx, my = event.pos
                # 收集阳光
                for s in sun_items[:]:
                    dist = ((mx - s.x)**2 + (my - s.y)**2)**0.5
                    if dist < s.radius + 10:
                        sun_count += 25
                        sun_items.remove(s)
                # 种植植物
                grid_x = mx // GRID_SIZE * GRID_SIZE
                grid_y = my // GRID_SIZE * GRID_SIZE
                # 不种在UI区域
                if grid_y >= ROW_COUNT * GRID_SIZE:
                    continue
                # 判断格子是否已有植物
                has_plant = False
                for p in plants:
                    if p.x == grid_x and p.y == grid_y:
                        has_plant = True
                if has_plant:
                    continue
                # 种植向日葵
                if select_plant == 0 and sun_count >= SUN_COST_SUNFLOWER:
                    plants.append(SunFlower(grid_x, grid_y))
                    sun_count -= SUN_COST_SUNFLOWER
                # 种植豌豆射手
                elif select_plant == 1 and sun_count >= SUN_COST_PEASHOOTER:
                    plants.append(PeaShooter(grid_x, grid_y))
                    sun_count -= SUN_COST_PEASHOOTER
            # 键盘切换植物
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_1:
                    select_plant = 0
                if event.key == pygame.K_2:
                    select_plant = 1

        if not game_over:
            # 1. 生成僵尸
            spawn_zombie_timer -= 1
            if spawn_zombie_timer <= 0:
                rand_row = random.randint(0, ROW_COUNT-1)
                zombies.append(Zombie(rand_row))
                spawn_zombie_timer = random.randint(200, 400)

            # 2. 更新阳光
            for s in sun_items[:]:
                s.life -= 1
                if s.life <= 0:
                    sun_items.remove(s)

            # 3. 更新植物（产阳光、射击）
            for p in plants:
                if isinstance(p, SunFlower):
                    p.update(sun_items)
                elif isinstance(p, PeaShooter):
                    p.update(bullets, zombies)

            # 4. 更新子弹 + 子弹碰撞僵尸
            for b in bullets[:]:
                b.update()
                if b.x > WIDTH:
                    bullets.remove(b)
                    continue
                # 碰撞僵尸
                for z in zombies[:]:
                    dist = ((b.x - (z.x+40))**2 + (b.y - (z.y+40))**2)**0.5
                    if dist < b.r + 40:
                        z.hp -= b.damage
                        bullets.remove(b)
                        if z.hp <= 0:
                            zombies.remove(z)
                        break

            # 5. 更新僵尸，检测游戏失败
            for z in zombies[:]:
                z.update(plants)
                if z.x < 0:  # 僵尸走到最左边，游戏结束
                    game_over = True
                if z.hp <= 0:
                    zombies.remove(z)

            # 6. 删除被啃死的植物
            for p in plants[:]:
                if p.hp <= 0:
                    plants.remove(p)

        # ===================== 绘制所有元素 =====================
        # 阳光道具
        for s in sun_items:
            s.draw()
        # 植物
        for p in plants:
            p.draw()
        # 子弹
        for b in bullets:
            b.draw()
        # 僵尸
        for z in zombies:
            z.draw()

        # UI文字（英文显示，适配Arial字体）
        sun_text = font.render(f"Sun: {sun_count}", True, WHITE)
        screen.blit(sun_text, (10, HEIGHT-40))
        plant_text = font.render(f"[1]SunFlower 50 | [2]Peashooter 100", True, WHITE)
        screen.blit(plant_text, (200, HEIGHT-40))

        # 游戏结束画面
        if game_over:
            over_text = font.render("Game Over! Zombie ate your brain", True, RED)
            screen.blit(over_text, (WIDTH//2 - 240, HEIGHT//2))

        pygame.display.flip()
        clock.tick(FPS)

if __name__ == "__main__":
    main()