import pygame
import random
import sys

pygame.init()
SCREEN_WIDTH = 900
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("植物大战僵尸（寒冰射手版）")
clock = pygame.time.Clock()
FPS = 60

WHITE = (255, 255, 255)
GREEN = (80, 180, 60)
BROWN = (120, 90, 50)
RED = (220, 30, 30)
BLUE = (50, 120, 220)
ICE_BLUE = (100, 200, 255)
BLACK = (0, 0, 0)
GRAY = (150, 150, 150)
ORANGE = (255,120,0)
YELLOW = (255,255,0)

ROW_COUNT = 5
COL_COUNT = 8
GRID_W = 80
GRID_H = 100
GRID_START_X = 100
GRID_START_Y = 80

# 兜底字体，稳定不报错
try:
    font = pygame.font.Font(None, 24)
    small_font = pygame.font.Font(None, 18)
except Exception:
    font = pygame.font.Font(None, 24)
    small_font = pygame.font.Font(None, 18)


class Plant(pygame.sprite.Sprite):
    def __init__(self, row, col, plant_type="peashooter"):
        super().__init__()
        self.row = row
        self.col = col
        self.x = GRID_START_X + col * GRID_W + GRID_W // 2
        self.y = GRID_START_Y + row * GRID_H + GRID_H // 2
        self.type = plant_type
        self.hp = 100
        self.attack_cd = 0
        self.attack_interval = 90
        self.radius = 32
        self.exploding = False
        self.explode_radius = 0

        if self.type == "peashooter":
            self.color = GREEN
            self.cost = 100
        elif self.type == "iceshooter":
            self.color = ICE_BLUE
            self.cost = 175
        elif self.type == "sunflower":
            self.color = (255, 220, 0)
            self.cost = 50
            self.attack_interval = 180
        elif self.type == "wallnut":
            self.color = (139, 90, 43)
            self.cost = 50
            self.hp = 400
        elif self.type == "cherrybomb":
            self.color = (200, 30, 30)
            self.cost = 150
            self.fuse = 60

    def update(self, sun_list, zombie_group):
        self.attack_cd -= 1
        if self.type == "sunflower":
            if self.attack_cd <= 0:
                sun_list.append(Sun(self.x, self.y - 30))
                self.attack_cd = self.attack_interval
        elif self.type == "peashooter":
            if self.attack_cd <= 0:
                for z in zombie_group:
                    if z.row == self.row:
                        bullet = Bullet(self.x + 25, self.y, is_ice=False)
                        all_bullets.add(bullet)
                        self.attack_cd = self.attack_interval
                        break
        elif self.type == "iceshooter":
            if self.attack_cd <= 0:
                for z in zombie_group:
                    if z.row == self.row:
                        bullet = Bullet(self.x + 25, self.y, is_ice=True)
                        all_bullets.add(bullet)
                        self.attack_cd = self.attack_interval
                        break
        elif self.type == "cherrybomb":
            self.fuse -= 1
            if self.fuse <= 0 and not self.exploding:
                self.exploding = True
                for z in zombie_group:
                    if abs(z.row - self.row) <= 1 and abs(z.x - self.x) < 130:
                        z.hp = 0
            if self.exploding:
                self.explode_radius += 8
                if self.explode_radius > 140:
                    self.hp = 0

    def draw(self, surf):
        if self.exploding:
            pygame.draw.circle(surf, ORANGE, (int(self.x), int(self.y)), self.explode_radius, 4)
            pygame.draw.circle(surf, YELLOW, (int(self.x), int(self.y)), self.explode_radius//2, 2)
        else:
            pygame.draw.circle(surf, self.color, (int(self.x), int(self.y)), self.radius)
            if self.type == "peashooter":
                pygame.draw.circle(surf, (30, 100, 30), (int(self.x), int(self.y)), self.radius - 12)
            elif self.type == "iceshooter":
                pygame.draw.circle(surf, (40,140,200), (int(self.x), int(self.y)), self.radius - 12)
            elif self.type == "cherrybomb":
                pygame.draw.circle(surf, BLACK, (int(self.x)+12, int(self.y)-10),6)


class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y, is_ice=False):
        super().__init__()
        self.x = x
        self.y = y
        self.speed = 6
        self.damage = 25
        self.radius = 8
        self.is_ice = is_ice
        if self.is_ice:
            self.color = ICE_BLUE
        else:
            self.color = (80, 160, 255)

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

    def draw(self, surf):
        pygame.draw.circle(surf, self.color, (int(self.x), int(self.y)), self.radius)


class Zombie(pygame.sprite.Sprite):
    def __init__(self, row, z_type="normal"):
        super().__init__()
        self.row = row
        self.col = COL_COUNT
        self.x = GRID_START_X + self.col * GRID_W + GRID_W // 2
        self.y = GRID_START_Y + row * GRID_H + GRID_H // 2
        self.z_type = z_type
        if z_type == "normal":
            self.hp = 100
            self.color = (100, 130, 80)
        elif z_type == "cone":
            self.hp = 220
            self.color = (220, 140, 0)
        elif z_type == "bucket":
            self.hp = 500
            self.color = (90, 90, 90)
        self.base_speed = 0.4
        self.speed = self.base_speed
        self.slow_timer = 0
        self.attack_cd = 0
        self.radius = 36
        self.eating_plant = None

    def update(self, plant_group):
        # 减速倒计时
        if self.slow_timer > 0:
            self.slow_timer -=1
            self.speed = self.base_speed * 0.4
        else:
            self.speed = self.base_speed

        self.attack_cd -= 1
        self.eating_plant = None
        for p in plant_group:
            if p.row == self.row and abs(p.x - self.x) < 45:
                self.eating_plant = p
                break
        if self.eating_plant:
            if self.attack_cd <= 0:
                self.eating_plant.hp -= 10
                self.attack_cd = 40
        else:
            self.x -= self.speed
        if self.x < 60:
            global game_over
            game_over = True

    def draw(self, surf):
        # 被减速的时候轮廓变蓝提示
        pygame.draw.circle(surf, self.color, (int(self.x), int(self.y)), self.radius)
        if self.slow_timer > 0:
            pygame.draw.circle(surf, ICE_BLUE, (int(self.x), int(self.y)), self.radius+3,2)
        pygame.draw.circle(surf, RED, (int(self.x) - 12, int(self.y) - 8), 5)
        pygame.draw.circle(surf, RED, (int(self.x) + 12, int(self.y) - 8), 5)


class Sun(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.x = x
        self.y = y
        self.vy = 2
        self.life = 300
        self.radius = 16

    def update(self):
        self.life -= 1
        if self.y < GRID_START_Y + GRID_H * ROW_COUNT:
            self.y += self.vy
        if self.life <= 0:
            self.kill()

    def draw(self, surf):
        pygame.draw.circle(surf, (255, 210, 0), (int(self.x), int(self.y)), self.radius)


class Button:
    def __init__(self, x, y, w, h, text, cost, plant_type):
        self.rect = pygame.Rect(x, y, w, h)
        self.text = text
        self.cost = cost
        self.plant_type = plant_type

    def draw(self, surf, sun_num):
        color = GRAY if sun_num < self.cost else (200, 200, 200)
        pygame.draw.rect(surf, color, self.rect)
        pygame.draw.rect(surf, BLACK, self.rect, 2)
        t1 = small_font.render(self.text, True, BLACK)
        t2 = small_font.render(f"{self.cost}", True, RED)
        surf.blit(t1, (self.rect.x + 5, self.rect.y + 3))
        surf.blit(t2, (self.rect.x + 5, self.rect.y + 24))

    def is_click(self, pos):
        return self.rect.collidepoint(pos)


all_plants = pygame.sprite.Group()
all_zombies = pygame.sprite.Group()
all_bullets = pygame.sprite.Group()
all_suns = pygame.sprite.Group()

sun_count = 100000000000000000000000000000000000000000000000000
selected_plant = None
game_over = False
spawn_zombie_timer = 0

btn1 = Button(30, 20, 85, 50, "豌豆", 100, "peashooter")
btn2 = Button(130, 20, 85, 50, "向日葵", 50, "sunflower")
btn3 = Button(230, 20, 85, 50, "坚果", 50, "wallnut")
btn4 = Button(330, 20, 85, 50, "樱桃", 150, "cherrybomb")
btn5 = Button(430, 20, 85, 50, "寒冰", 175, "iceshooter")
button_list = [btn1, btn2, btn3, btn4, btn5]


def main():
    global sun_count, selected_plant, game_over, spawn_zombie_timer
    running = True
    while running:
        clock.tick(FPS)
        screen.fill((40, 80, 30))
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
                mx, my = event.pos
                for s in all_suns:
                    if ((mx - s.x)**2 + (my - s.y)**2) < s.radius**2:
                        sun_count += 25
                        s.kill()
                        break
                else:
                    clicked_btn = None
                    for b in button_list:
                        if b.is_click((mx, my)):
                            clicked_btn = b
                            break
                    if clicked_btn:
                        if sun_count >= clicked_btn.cost:
                            selected_plant = clicked_btn
                    else:
                        if selected_plant:
                            col = (mx - GRID_START_X) // GRID_W
                            row = (my - GRID_START_Y) // GRID_H
                            if 0 <= row < ROW_COUNT and 0 <= col < COL_COUNT:
                                occupied = False
                                for p in all_plants:
                                    if p.row == row and p.col == col:
                                        occupied = True
                                        break
                                if not occupied:
                                    new_p = Plant(row, col, selected_plant.plant_type)
                                    all_plants.add(new_p)
                                    sun_count -= selected_plant.cost
                                    selected_plant = None

        if not game_over:
            spawn_zombie_timer += 1
            if spawn_zombie_timer > 300:
                r = random.randint(0, ROW_COUNT - 1)
                rd = random.random()
                if rd < 0.5:
                    z = Zombie(r, "normal")
                elif rd < 0.8:
                    z = Zombie(r, "cone")
                else:
                    z = Zombie(r, "bucket")
                all_zombies.add(z)
                spawn_zombie_timer = 0

            for p in all_plants:
                p.update(all_suns, all_zombies)
            for z in all_zombies:
                z.update(all_plants)
            for b in all_bullets:
                b.update()
            for s in all_suns:
                s.update()

            # 子弹碰撞
            for bullet in all_bullets:
                for z in all_zombies:
                    dist = ((bullet.x - z.x)**2 + (bullet.y - z.y)**2)**0.5
                    if dist < z.radius:
                        z.hp -= bullet.damage
                        if bullet.is_ice:
                            z.slow_timer = 90 # 减速90帧
                        bullet.kill()
                        if z.hp <= 0:
                            z.kill()
                        break
            for p in all_plants:
                if p.hp <= 0:
                    p.kill()

        for r in range(ROW_COUNT):
            for c in range(COL_COUNT):
                rect = pygame.Rect(GRID_START_X + c * GRID_W, GRID_START_Y + r * GRID_H, GRID_W - 2, GRID_H - 2)
                pygame.draw.rect(screen, (100, 170, 60), rect)
                pygame.draw.rect(screen, (60, 130, 30), rect, 1)

        for b in button_list:
            b.draw(screen, sun_count)

        sun_text = font.render(f"阳光：{sun_count}", True, WHITE)
        screen.blit(sun_text, (540, 25))

        if selected_plant:
            tip = small_font.render(f"已选择：{selected_plant.text}，点击草地放置", True, WHITE)
            screen.blit(tip, (650, 25))

        for s in all_suns:
            s.draw(screen)
        for p in all_plants:
            p.draw(screen)
        for z in all_zombies:
            z.draw(screen)
        for b in all_bullets:
            b.draw(screen)

        if game_over:
            over_text = font.render("游戏结束！僵尸吃掉你的脑子！", True, RED)
            screen.blit(over_text, (280, 260))

        pygame.display.flip()
    pygame.quit()
    sys.exit()


if __name__ == "__main__":
    main()