import pygame
import random
import sys

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

# 颜色
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (34, 139, 34)
LIGHT_GREEN = (102, 204, 102)
BROWN = (139, 69, 19)
SKY_BLUE = (135, 206, 235)
YELLOW = (255, 215, 0)
RED = (220, 30, 30)
GRAY = (100, 100, 100)

# 网格设置
GRID_W = 80
GRID_H = 80
ROW_NUM = 5
COL_NUM = 9
OFFSET_X = 50
OFFSET_Y = 80
BAR_HEIGHT = 70

# 游戏变量
sun = 150
selected_plant = None
plant_cost = {"peashooter": 100}
game_over = False

# 【修复字体】改用 pygame.font.Font 加载内置默认字体，避开SysFont bug
try:
    font36 = pygame.font.Font(pygame.font.get_default_font(), 36)
    font72 = pygame.font.Font(pygame.font.get_default_font(), 72)
except:
    font36 = pygame.font.SysFont("arial", 36)
    font72 = pygame.font.SysFont("arial", 72)

# 精灵组
plant_group = pygame.sprite.Group()
bullet_group = pygame.sprite.Group()
zombie_group = pygame.sprite.Group()

# 植物
class Plant(pygame.sprite.Sprite):
    def __init__(self, x, y, ptype):
        super().__init__()
        self.type = ptype
        self.x = x
        self.y = y
        self.rect = pygame.Rect(x, y, GRID_W - 5, GRID_H - 5)
        self.hp = 100
        self.shoot_timer = 0
        self.shoot_cd = 90

    def update(self):
        if self.type == "peashooter":
            self.shoot_timer += 1
            if self.shoot_timer >= self.shoot_cd:
                self.shoot_timer = 0
                bullet = Bullet(self.rect.centerx + 20, self.rect.centery)
                bullet_group.add(bullet)

    def draw(self):
        pygame.draw.rect(screen, GREEN, self.rect, border_radius=8)
        pygame.draw.circle(screen, LIGHT_GREEN, self.rect.center, 25)

# 子弹
class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.x = x
        self.y = y
        self.speed = 6
        self.radius = 8
        self.rect = pygame.Rect(x - self.radius, y - self.radius, self.radius * 2, self.radius * 2)

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

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

# 僵尸
class Zombie(pygame.sprite.Sprite):
    def __init__(self, row):
        super().__init__()
        self.row = row
        self.x = WIDTH - 20
        self.y = OFFSET_Y + row * GRID_H
        self.speed = 0.6
        self.hp = 80
        self.rect = pygame.Rect(self.x, self.y, GRID_W - 10, GRID_H - 10)

    def update(self):
        self.x -= self.speed
        self.rect.x = self.x
        if self.x < 0:
            global game_over
            game_over = True

    def draw(self):
        pygame.draw.rect(screen, GRAY, self.rect, border_radius=6)
        pygame.draw.circle(screen, RED, (self.rect.centerx, self.rect.top + 25), 12)

# UI按钮
peashooter_btn = pygame.Rect(20, 10, 60, 50)

def draw_ui():
    pygame.draw.rect(screen, BROWN, (0, 0, WIDTH, BAR_HEIGHT))
    sun_text = font36.render(f"阳光: {sun}", True, YELLOW)
    screen.blit(sun_text, (100, 15))
    btn_color = LIGHT_GREEN if selected_plant == "peashooter" else GREEN
    pygame.draw.rect(screen, btn_color, peashooter_btn, border_radius=5)
    cost_text = font36.render("100", True, WHITE)
    screen.blit(cost_text, (30, 18))

def draw_grass():
    for r in range(ROW_NUM):
        for c in range(COL_NUM):
            x = OFFSET_X + c * GRID_W
            y = OFFSET_Y + r * GRID_H
            rect = pygame.Rect(x, y, GRID_W, GRID_H)
            if (r + c) % 2 == 0:
                clr = (80, 160, 60)
            else:
                clr = (60, 130, 40)
            pygame.draw.rect(screen, clr, rect)
            pygame.draw.rect(screen, BLACK, rect, 1)

def get_grid_pos(mx, my):
    if my < OFFSET_Y or mx < OFFSET_X:
        return None
    c = (mx - OFFSET_X) // GRID_W
    r = (my - OFFSET_Y) // GRID_H
    if 0 <= r < ROW_NUM and 0 <= c < COL_NUM:
        return (r, c)
    return None

sun_timer = 0
zombie_spawn_timer = 0

running = True
while running:
    clock.tick(FPS)
    screen.fill(SKY_BLUE)

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
            mx, my = pygame.mouse.get_pos()
            # 点击植物按钮
            if peashooter_btn.collidepoint(mx, my):
                selected_plant = "peashooter"
            # 放置植物
            grid = get_grid_pos(mx, my)
            if grid and selected_plant:
                r, c = grid
                px = OFFSET_X + c * GRID_W
                py = OFFSET_Y + r * GRID_H
                occupied = False
                for p in plant_group:
                    if p.rect.x == px and p.rect.y == py:
                        occupied = True
                        break
                if not occupied and sun >= plant_cost[selected_plant]:
                    sun -= plant_cost[selected_plant]
                    plant_group.add(Plant(px, py, selected_plant))

    if not game_over:
        sun_timer += 1
        if sun_timer >= 240:
            sun += 25
            sun_timer = 0

        zombie_spawn_timer += 1
        if zombie_spawn_timer >= 400:
            r = random.randint(0, ROW_NUM - 1)
            zombie_group.add(Zombie(r))
            zombie_spawn_timer = 0

        plant_group.update()
        bullet_group.update()
        zombie_group.update()

        # 子弹击中僵尸
        collision = pygame.sprite.groupcollide(zombie_group, bullet_group, False, True)
        for zom, bull_list in collision.items():
            zom.hp -= 25
            if zom.hp <= 0:
                zom.kill()

    draw_grass()
    draw_ui()

    for p in plant_group:
        p.draw()
    for b in bullet_group:
        b.draw()
    for z in zombie_group:
        z.draw()

    if game_over:
        over_surf = font72.render("游戏结束！", True, RED)
        screen.blit(over_surf, (WIDTH // 2 - 150, HEIGHT // 2))

    pygame.display.flip()

pygame.quit()
sys.exit()