import pygame
import random
import sys

# 初始化Pygame
pygame.init()
WIDTH, HEIGHT = 900, 600
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("植物大战僵尸 Python版")
clock = pygame.time.Clock()

# 颜色常量
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (34, 139, 34)
BROWN = (139, 69, 19)
SKY_BLUE = (135, 206, 235)
YELLOW = (255, 255, 0)
RED = (255, 0, 0)

# 网格设置（草坪6行9列）
ROW_NUM = 6
COL_NUM = 9
CELL_W = 80
CELL_H = 80
OFFSET_X = 50
OFFSET_Y = 80

# 阳光数值
sun = 50
font = pygame.font.SysFont("simhei", 28)
tip_font = pygame.font.SysFont("simhei", 20)

# 植物类型：向日葵、豌豆射手
PLANT_SUNFLOWER = 1
PLANT_PEASHOOTER = 2
# 选中要种植的植物
select_plant = None

# 存放所有植物、僵尸、子弹、阳光掉落物
plants = []
zombies = []
bullets = []
sun_items = []

# 植物基类


class Plant:
    def __init__(self, row, col, ptype):
        self.row = row
        self.col = col
        self.x = OFFSET_X + col * CELL_W
        self.y = OFFSET_Y + row * CELL_H
        self.type = ptype
        self.hp = 100
        self.cd = 0  # 冷却计时器

    def update(self):
        self.cd += 1
        if self.type == PLANT_SUNFLOWER:
            # 向日葵每120帧产出阳光
            if self.cd >= 120:
                sun_items.append(SunItem(self.x + 30, self.y))
                self.cd = 0
        elif self.type == PLANT_PEASHOOTER:
            # 豌豆射手前方有僵尸就发射子弹
            if self.cd >= 60:
                has_enemy = False
                for z in zombies:
                    if z.row == self.row and z.x > self.x:
                        has_enemy = True
                        break
                if has_enemy:
                    bullets.append(Bullet(self.x + 60, self.y + 30, self.row))
                    self.cd = 0

    def draw(self):
        if self.type == PLANT_SUNFLOWER:
            pygame.draw.circle(
                screen, YELLOW, (int(self.x+40), int(self.y+40)), 35)
            pygame.draw.circle(
                screen, BROWN, (int(self.x+40), int(self.y+40)), 15)
        elif self.type == PLANT_PEASHOOTER:
            pygame.draw.rect(screen, GREEN, (self.x+10, self.y+10, 60, 60))
            pygame.draw.circle(screen, (0, 180, 0),
                               (int(self.x+70), int(self.y+40)), 12)
        # 血条
        pygame.draw.rect(screen, RED, (self.x, self.y-8, 70, 4))
        pygame.draw.rect(screen, GREEN, (self.x, self.y-8, self.hp*0.7, 4))

# 子弹类


class Bullet:
    def __init__(self, x, y, row):
        self.x = x
        self.y = y
        self.row = row
        self.speed = 5
        self.dmg = 25

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

    def draw(self):
        pygame.draw.circle(screen, (50, 200, 50),
                           (int(self.x), int(self.y)), 8)

# 僵尸类


class Zombie:
    def __init__(self):
        self.row = random.randint(0, ROW_NUM-1)
        self.x = WIDTH
        self.y = OFFSET_Y + self.row * CELL_H
        self.hp = 150
        self.speed = 0.3
        self.eat_cd = 0

    def update(self):
        # 检测前方有没有植物
        hit_plant = None
        for p in plants:
            if p.row == self.row and abs(p.x - self.x) < 40:
                hit_plant = p
                break
        if hit_plant:
            # 啃植物
            self.eat_cd += 1
            if self.eat_cd >= 30:
                hit_plant.hp -= 10
                self.eat_cd = 0
        else:
            self.x -= self.speed

    def draw(self):
        # 全部坐标强制转int，修复float报错
        ix = int(self.x)
        iy = int(self.y)
        pygame.draw.rect(screen, (120, 100, 80), (ix, iy, 65, 75))
        pygame.draw.circle(screen, WHITE, (ix+15, iy+20), 7)
        pygame.draw.circle(screen, WHITE, (ix+45, iy+20), 7)
        # 血条
        pygame.draw.rect(screen, RED, (ix, iy-6, 65, 4))
        pygame.draw.rect(screen, GREEN, (ix, iy-6, self.hp*0.43, 4))

# 掉落阳光


class SunItem:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.timer = 300  # 存在时间

    def update(self):
        self.timer -= 1

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

# 绘制草坪背景


def draw_grass():
    screen.fill(SKY_BLUE)
    # 绘制草地大背景
    pygame.draw.rect(screen, GREEN, (OFFSET_X, OFFSET_Y,
                     COL_NUM*CELL_W, ROW_NUM*CELL_H))
    # 网格线
    for r in range(ROW_NUM+1):
        y = OFFSET_Y + r * CELL_H
        pygame.draw.line(screen, BLACK, (OFFSET_X, y),
                         (OFFSET_X+COL_NUM*CELL_W, y), 2)
    for c in range(COL_NUM+1):
        x = OFFSET_X + c * CELL_W
        pygame.draw.line(screen, BLACK, (x, OFFSET_Y),
                         (x, OFFSET_Y+ROW_NUM*CELL_H), 2)

# 顶部选卡栏绘制


def draw_card_bar():
    # 向日葵卡片
    pygame.draw.rect(screen, YELLOW, (20, 10, 80, 50))
    s_text = tip_font.render("向日葵50", True, BLACK)
    screen.blit(s_text, (22, 15))
    # 豌豆射手卡片
    pygame.draw.rect(screen, GREEN, (110, 10, 80, 50))
    p_text = tip_font.render("豌豆100", True, BLACK)
    screen.blit(p_text, (112, 15))
    # 阳光显示
    sun_text = font.render(f"阳光: {sun}", True, YELLOW)
    screen.blit(sun_text, (220, 15))


# 主游戏循环
spawn_zombie_timer = 0
game_over = False
win = False

while True:
    clock.tick(60)
    draw_grass()
    draw_card_bar()

    # 事件监听
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            pygame.quit()
            sys.exit()
        if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
            mx, my = pygame.mouse.get_pos()
            # 点击选植物卡片
            if 20 <= mx <= 100 and 10 <= my <= 60:
                select_plant = PLANT_SUNFLOWER
            if 110 <= mx <= 190 and 10 <= my <= 60:
                select_plant = PLANT_PEASHOOTER
            # 点击草坪种植
            if OFFSET_X <= mx <= OFFSET_X + COL_NUM*CELL_W and OFFSET_Y <= my <= OFFSET_Y + ROW_NUM*CELL_H:
                col = (mx - OFFSET_X) // CELL_W
                row = (my - OFFSET_Y) // CELL_H
                # 判断格子是否已有植物
                occupied = False
                for p in plants:
                    if p.row == row and p.col == col:
                        occupied = True
                        break
                if not occupied and select_plant is not None:
                    cost = 50 if select_plant == PLANT_SUNFLOWER else 100
                    if sun >= cost:
                        sun -= cost
                        plants.append(Plant(row, col, select_plant))
            # 点击拾取阳光
            for s in sun_items[:]:
                if ((mx - s.x)**2 + (my - s.y)**2) < 30**2:
                    sun += 25
                    sun_items.remove(s)

    if game_over:
        over_text = font.render("游戏结束！僵尸吃掉脑子啦", True, RED)
        screen.blit(over_text, (WIDTH//2-180, HEIGHT//2))
        pygame.display.update()
        continue

    # 定时生成僵尸
    spawn_zombie_timer += 1
    if spawn_zombie_timer >= 400:
        zombies.append(Zombie())
        spawn_zombie_timer = 0

    # 更新&绘制植物
    for p in plants[:]:
        p.update()
        p.draw()
        if p.hp <= 0:
            plants.remove(p)

    # 更新&绘制子弹 + 子弹打僵尸
    for b in bullets[:]:
        b.update()
        b.draw()
        if b.x > WIDTH:
            bullets.remove(b)
            continue
        for z in zombies[:]:
            if abs(b.x - z.x) < 35 and abs(b.y - z.y) < 40:
                z.hp -= b.dmg
                bullets.remove(b)
                break

    # 更新&绘制僵尸
    for z in zombies[:]:
        z.update()
        z.draw()
        # 僵尸走到最左边，游戏失败
        if z.x < OFFSET_X - 20:
            game_over = True
        if z.hp <= 0:
            zombies.remove(z)

    # 更新掉落阳光
    for s in sun_items[:]:
        s.update()
        s.draw()
        if s.timer <= 0:
            sun_items.remove(s)

    # 胜利判定：坚持一定时间僵尸不再刷新（简化版）
    if len(zombies) == 0 and spawn_zombie_timer > 1200 and len(plants) > 2:
        win_text = font.render("恭喜通关！", True, (0, 255, 0))
        screen.blit(win_text, (WIDTH//2-100, HEIGHT//2))

    pygame.display.flip()