import pygame
import random
import sys

# ===================== 游戏初始化 =====================
pygame.init()
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 500
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("植物大战僵尸 - 含防御坚果+爆炸植物")
clock = pygame.time.Clock()
FPS = 60

# 颜色定义
GREEN = (0, 255, 0)
BROWN = (139, 69, 19)
YELLOW = (255, 255, 0)
RED = (255, 0, 0)
BLUE = (0, 0, 255)
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
ORANGE = (255, 120, 0)
GRAY = (120, 120, 120)

# 游戏全局数据
sunlight = 150
game_over = False
win = False

GRID_X = 150
GRID_Y = 80
CELL_SIZE = 70

# ===================== 游戏类定义 =====================
class Sun:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.value = 25
        self.radius = 15
        self.speed = 1

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

    def fall(self):
        self.y += self.speed

class Plant:
    def __init__(self, x, y, hp, cost):
        self.x = x
        self.y = y
        self.hp = hp
        self.cost = cost
        self.rect = pygame.Rect(x, y, CELL_SIZE-5, CELL_SIZE-5)

# 向日葵 30阳光
class Sunflower(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 100, 30)
        self.sun_timer = 0

    def draw(self):
        pygame.draw.rect(screen, (255, 200, 0), self.rect)

    def produce_sun(self):
        self.sun_timer += 1
        if self.sun_timer >= 180:
            self.sun_timer = 0
            return Sun(self.x + CELL_SIZE//2, self.y - 10)
        return None

# 豌豆射手 60阳光
class Peashooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 100, 60)
        self.shoot_timer = 0

    def draw(self):
        pygame.draw.rect(screen, GREEN, self.rect)

    def shoot(self):
        self.shoot_timer += 1
        if self.shoot_timer >= 45:
            self.shoot_timer = 0
            return Bullet(self.x + CELL_SIZE, self.y + CELL_SIZE//2)
        return None

# 防御坚果【新增】高血量纯防御
class WallNut(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, 800, 40)  # 超高血量，便宜
    def draw(self):
        pygame.draw.rect(screen, GRAY, self.rect, border_radius=12)

# 樱桃爆炸炸弹 90阳光
class CherryBomb(Plant):
    def __init__(self, x, y, row):
        super().__init__(x, y, 1, 90)
        self.row = row
        self.countdown = 60
        self.exploded = False

    def update(self, zombies):
        if self.exploded:
            return
        self.countdown -= 1
        if self.countdown <= 0:
            self.explode(zombies)

    def explode(self, zombies):
        self.exploded = True
        for z in zombies:
            if abs(z.y - self.y) < CELL_SIZE:
                z.hp -= 999

    def draw(self):
        color = RED if self.countdown > 20 else ORANGE
        pygame.draw.rect(screen, color, self.rect, border_radius=8)
        font = pygame.font.SysFont(None, 20)
        t = font.render("炸弹", True, WHITE)
        screen.blit(t, (self.x+8, self.y+20))

# 爆炸特效
class Explosion:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.radius = 10
        self.max_radius = 180
        self.alpha = 255

    def update(self):
        self.radius += 3
        self.alpha -= 8

    def draw(self):
        surface = pygame.Surface((self.max_radius*2, self.max_radius*2), pygame.SRCALPHA)
        color = (*(255,80,0), self.alpha)
        pygame.draw.circle(surface, color, (self.max_radius, self.max_radius), self.radius)
        screen.blit(surface, (self.x - self.max_radius, self.y - self.max_radius))

class Bullet:
    def __init__(self, x, y):
        self.x = x
        self.y = y
        self.speed = 5
        self.damage = 25
        self.radius = 8

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

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

class Zombie:
    def __init__(self, row):
        self.x = SCREEN_WIDTH - 50
        self.y = GRID_Y + row * CELL_SIZE + 5
        self.hp = 200
        self.speed = 0.5
        self.rect = pygame.Rect(self.x, self.y, CELL_SIZE, CELL_SIZE)

    def draw(self):
        pygame.draw.rect(screen, BROWN, self.rect)
        pygame.draw.rect(screen, RED, (self.x, self.y-10, CELL_SIZE, 8))
        pygame.draw.rect(screen, GREEN, (self.x, self.y-10, CELL_SIZE*(self.hp/200), 8))

    def move(self):
        self.x -= self.speed
        self.rect.x = self.x

# ===================== 游戏对象列表 =====================
plants = []
sun_list = []
bullets = []
zombies = []
explosions = []
selected_plant = None

zombie_spawn_timer = 0
zombie_spawn_interval = 300

# ===================== 绘制UI =====================
def draw_background():
    screen.fill((30, 30, 30))
    for row in range(5):
        for col in range(9):
            x = GRID_X + col * CELL_SIZE
            y = GRID_Y + row * CELL_SIZE
            pygame.draw.rect(screen, (50, 150, 50), (x, y, CELL_SIZE, CELL_SIZE), 2)

def draw_ui():
    font = pygame.font.SysFont(None, 40)
    sun_text = font.render(f"阳光: {sunlight}", True, YELLOW)
    screen.blit(sun_text, (20, 20))

    pygame.draw.rect(screen, (100, 100, 100), (20, 80, 100, 380))
    # 向日葵 30
    pygame.draw.rect(screen, (255, 200, 0), (30, 100, 80, 55))
    t1 = pygame.font.SysFont(None, 20).render("向日葵 30", True, BLACK)
    screen.blit(t1, (35, 115))
    # 豌豆 60
    pygame.draw.rect(screen, GREEN, (30, 165, 80, 55))
    t2 = pygame.font.SysFont(None, 20).render("豌豆 60", True, BLACK)
    screen.blit(t2, (35, 180))
    # 坚果防御 40
    pygame.draw.rect(screen, GRAY, (30, 230, 80, 55))
    t4 = pygame.font.SysFont(None, 20).render("坚果 40", True, WHITE)
    screen.blit(t4, (35, 245))
    # 樱桃炸弹 90
    pygame.draw.rect(screen, RED, (30, 295, 80, 55))
    t3 = pygame.font.SysFont(None, 20).render("炸弹 90", True, WHITE)
    screen.blit(t3, (35, 310))

def draw_game_over():
    font = pygame.font.SysFont(None, 60)
    text = font.render("僵尸吃掉脑子！", True, RED) if not win else font.render("胜利！", True, GREEN)
    screen.blit(text, (SCREEN_WIDTH//2 - 180, SCREEN_HEIGHT//2))

# ===================== 主循环 =====================
running = True
while running:
    clock.tick(FPS)
    draw_background()
    draw_ui()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

        if not game_over and event.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()
            # 选择植物
            if 30 <= mx <= 110:
                if 100 <= my <= 155 and sunlight >= 30:
                    selected_plant = "sunflower"
                elif 165 <= my <= 220 and sunlight >= 60:
                    selected_plant = "peashooter"
                elif 230 <= my <= 285 and sunlight >= 40:
                    selected_plant = "wallnut"
                elif 295 <= my <= 350 and sunlight >= 90:
                    selected_plant = "cherry"
            # 种植
            else:
                col = (mx - GRID_X) // CELL_SIZE
                row = (my - GRID_Y) // CELL_SIZE
                if 0 <= row < 5 and 0 <= col < 9:
                    x = GRID_X + col * CELL_SIZE
                    y = GRID_Y + row * CELL_SIZE
                    can_place = True
                    for p in plants:
                        if abs(p.x - x) < 10 and abs(p.y - y) < 10:
                            can_place = False
                            break
                    if can_place and selected_plant:
                        if selected_plant == "sunflower":
                            plants.append(Sunflower(x, y))
                            sunlight -= 30
                        elif selected_plant == "peashooter":
                            plants.append(Peashooter(x, y))
                            sunlight -= 60
                        elif selected_plant == "wallnut":
                            plants.append(WallNut(x, y))
                            sunlight -= 40
                        elif selected_plant == "cherry":
                            plants.append(CherryBomb(x, y, row))
                            sunlight -= 90
                        selected_plant = None
            # 捡阳光
            for sun in sun_list[:]:
                if (mx - sun.x)**2 + (my - sun.y)**2 <= sun.radius**2:
                    sunlight += sun.value
                    sun_list.remove(sun)

    if game_over:
        draw_game_over()
        pygame.display.flip()
        continue

    # 生成僵尸
    zombie_spawn_timer += 1
    if zombie_spawn_timer >= zombie_spawn_interval:
        zombie_spawn_timer = 0
        zombies.append(Zombie(random.randint(0, 4)))
        if zombie_spawn_interval > 120:
            zombie_spawn_interval -= 8

    # 阳光下落
    for sun in sun_list[:]:
        sun.fall()
        sun.draw()
        if sun.y > SCREEN_HEIGHT:
            sun_list.remove(sun)

    # 更新植物 & 爆炸
    for plant in plants[:]:
        plant.draw()
        if isinstance(plant, Sunflower):
            s = plant.produce_sun()
            if s:
                sun_list.append(s)
        elif isinstance(plant, Peashooter):
            b = plant.shoot()
            if b:
                bullets.append(b)
        elif isinstance(plant, CherryBomb):
            plant.update(zombies)
            if plant.exploded:
                explosions.append(Explosion(plant.x+35, plant.y+35))
                plants.remove(plant)

    # 爆炸特效
    for exp in explosions[:]:
        exp.update()
        exp.draw()
        if exp.alpha <= 0:
            explosions.remove(exp)

    # 子弹
    for bullet in bullets[:]:
        bullet.move()
        bullet.draw()
        if bullet.x > SCREEN_WIDTH:
            bullets.remove(bullet)

    # 僵尸
    for zombie in zombies[:]:
        zombie.move()
        zombie.draw()
        # 啃所有植物（坚果专门用来抗伤害）
        for p in plants[:]:
            if zombie.rect.colliderect(p.rect) and not isinstance(p, CherryBomb):
                zombie.speed = 0
                p.hp -= 0.5
                if p.hp <= 0:
                    plants.remove(p)
                    zombie.speed = 0.5
        # 子弹伤害
        for b in bullets[:]:
            if zombie.rect.collidepoint(b.x, b.y):
                zombie.hp -= b.damage
                if b in bullets:
                    bullets.remove(b)
        if zombie.hp <= 0:
            zombies.remove(zombie)
        if zombie.x <= GRID_X:
            game_over = True

    pygame.display.flip()

pygame.quit()
sys.exit()