import pygame
import random

# -------------------------- 初始化设置 --------------------------
pygame.init()
WIDTH, HEIGHT = 800, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("简易植物大战僵尸")
clock = pygame.time.Clock()
FPS = 60

# 颜色定义
WHITE = (255, 255, 255)
GREEN = (34, 139, 34)
BROWN = (139, 69, 19)
RED = (255, 0, 0)
YELLOW = (255, 215, 0)
SKY_BLUE = (135, 206, 235)
GRAY = (80, 80, 80)

# 游戏参数
GRID_COLS = 9
GRID_ROWS = 5
CELL_W = WIDTH // GRID_COLS
CELL_H = (HEIGHT - 100) // GRID_ROWS
sun = 50  # 初始阳光

# -------------------------- 游戏对象类 --------------------------
class Plant:
    def __init__(self, x, y, plant_type):
        self.x = x
        self.y = y
        self.type = plant_type
        self.hp = 100
        self.shoot_timer = 0
        self.sun_timer = 0

    def update(self, bullets):
        if self.type == "peashooter":
            self.shoot_timer += 1
            if self.shoot_timer >= 90:
                bullets.append(Bullet(self.x + 40, self.y + 20))
                self.shoot_timer = 0
        elif self.type == "sunflower":
            self.sun_timer += 1

    def draw(self):
        if self.type == "peashooter":
            pygame.draw.circle(screen, GREEN, (self.x+25, self.y+25), 25)
        elif self.type == "sunflower":
            pygame.draw.circle(screen, YELLOW, (self.x+25, self.y+25), 25)

class Zombie:
    def __init__(self, row):
        self.y = row * CELL_H + 100
        self.x = WIDTH
        self.speed = 0.5
        self.hp = 100

    def update(self, plants):
        self.x -= self.speed
        # 啃植物检测
        for p in plants:
            if abs(p.y - self.y) < 40 and abs(p.x - self.x) < 40:
                self.x += self.speed
                p.hp -= 0.3

    def draw(self):
        pygame.draw.rect(screen, BROWN, (self.x, self.y, 45, 45))

class Bullet:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.speed = 4

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

    def draw(self):
        pygame.draw.circle(screen, (0, 200, 0), (int(self.x), int(self.y)), 6)

# -------------------------- 主游戏循环 --------------------------
def main():
    global sun
    plants = []
    zombies = []
    bullets = []
    spawn_zombie_timer = 0
    running = True
    selected_plant = None

    while running:
        screen.fill(SKY_BLUE)
        # 绘制草地网格
        for col in range(GRID_COLS):
            for row in range(GRID_ROWS):
                rect = pygame.Rect(col*CELL_W, row*CELL_H+100, CELL_W-2, CELL_H-2)
                pygame.draw.rect(screen, GREEN, rect)

        # 顶部灰色面板
        pygame.draw.rect(screen, GRAY, (0, 0, WIDTH, 100))

        # 两个选植物按钮
        sunflower_btn = pygame.Rect(150, 20, 80, 60)
        peashooter_btn = pygame.Rect(260, 20, 80, 60)
        pygame.draw.rect(screen, YELLOW, sunflower_btn)
        pygame.draw.rect(screen, GREEN, peashooter_btn)

        # 阳光圆形标识（黄色圆圈代表阳光数值）
        pygame.draw.circle(screen,YELLOW,(40,50),20)

        # 事件处理
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.MOUSEBUTTONDOWN:
                mx, my = pygame.mouse.get_pos()
                # 选植物：黄色按钮=向日葵，绿色按钮=豌豆射手
                if sunflower_btn.collidepoint(mx, my):
                    selected_plant = "sunflower"
                elif peashooter_btn.collidepoint(mx, my):
                    selected_plant = "peashooter"
                else:
                    # 在草地种植
                    if my > 100 and selected_plant:
                        col = mx // CELL_W
                        row = (my - 100) // CELL_H
                        px = col * CELL_W
                        py = row * CELL_H + 100
                        cost = 50 if selected_plant == "sunflower" else 100
                        # 判断格子是否已有植物
                        can_plant = True
                        for p in plants:
                            if p.x == px and p.y == py:
                                can_plant = False
                        if can_plant and sun >= cost:
                            plants.append(Plant(px, py, selected_plant))
                            sun -= cost
                            selected_plant = None

        # 向日葵产出阳光
        for p in plants[:]:
            if p.type == "sunflower" and p.sun_timer >= 180:
                sun += 25
                p.sun_timer = 0

        # 更新植物
        for p in plants:
            p.update(bullets)

        # 生成僵尸
        spawn_zombie_timer += 1
        if spawn_zombie_timer >= 300:
            r = random.randint(0, GRID_ROWS-1)
            zombies.append(Zombie(r * CELL_H + 100))
            spawn_zombie_timer = 0

        # 更新僵尸
        game_over = False
        for z in zombies[:]:
            z.update(plants)
            if z.x < 0:
                game_over = True
                print("游戏结束！僵尸吃掉了你！")
                break
        if game_over:
            running = False

        # 更新子弹+子弹打僵尸碰撞
        for b in bullets[:]:
            b.update()
            if b.x > WIDTH:
                bullets.remove(b)
                continue
            for z in zombies[:]:
                if abs(b.x - z.x) < 30 and abs(b.y - z.y) < 30:
                    z.hp -= 25
                    bullets.remove(b)
                    if z.hp <= 0:
                        zombies.remove(z)
                    break

        # 移除死掉的植物
        for p in plants[:]:
            if p.hp <= 0:
                plants.remove(p)

        # 绘制所有物体
        for p in plants:
            p.draw()
        for z in zombies:
            z.draw()
        for b in bullets:
            b.draw()

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

if __name__ == "__main__":
    main()