import pygame
import random
import sys

# ===================== 全局常量设置 =====================
WIDTH, HEIGHT = 900, 600
FPS = 60
GRID_W, GRID_H = 80, 100
ROW_NUM = 5
COL_NUM = 8

WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
GREEN = (34, 139, 34)
GREEN_LIGHT = (60, 160, 60)
BROWN = (139, 90, 43)
RED = (220, 0, 0)
YELLOW = (255, 215, 0)
SKY_BLUE = (135, 206, 235)
ICE_BLUE = (0, 190, 255)
DARK_GREEN = (0, 90, 0)
GRAY = (120, 120, 120)
ORANGE = (255, 140, 0)

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("植物大战僵尸 修复僵尸生成")
clock = pygame.time.Clock()

def get_font(size):
    try:
        return pygame.font.SysFont("Microsoft YaHei", size)
    except:
        try:
            return pygame.font.SysFont("SimHei", size)
        except:
            return pygame.font.Font(None, size)

font_big = get_font(40)
font_mid = get_font(28)
font_small = get_font(20)

# ===================== 植物父类 =====================
class Plant(pygame.sprite.Sprite):
    def __init__(self, x, y, cost, hp):
        super().__init__()
        self.rect = pygame.Rect(x, y, GRID_W - 10, GRID_H - 10)
        self.cost = cost
        self.hp = hp
        self.max_hp = hp
        self.row_index = y // GRID_H

# 向日葵
class SunFlower(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, cost=50, hp=100)
        self.produce_timer = 0
        self.produce_interval = 5000

    def update(self, dt, sun_group):
        self.produce_timer += dt
        if self.produce_timer >= self.produce_interval:
            self.produce_timer = 0
            sun_group.add(Sun(self.rect.centerx, self.rect.centery))

# 豌豆射手
class PeaShooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, cost=100, hp=100)
        self.shoot_timer = 0
        self.shoot_interval = 1200

    def update(self, dt, bullet_group, zombie_group):
        self.shoot_timer += dt
        have_zombie = any(z.row_index == self.row_index for z in zombie_group)
        if self.shoot_timer >= self.shoot_interval and have_zombie:
            self.shoot_timer = 0
            bullet_group.add(Bullet(self.rect.right, self.rect.centery, damage=25, speed=6))

# 寒冰射手
class IcePeaShooter(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, cost=175, hp=100)
        self.shoot_timer = 0
        self.shoot_interval = 1500

    def update(self, dt, bullet_group, zombie_group):
        self.shoot_timer += dt
        have_zombie = any(z.row_index == self.row_index for z in zombie_group)
        if self.shoot_timer >= self.shoot_interval and have_zombie:
            self.shoot_timer = 0
            bullet_group.add(IceBullet(self.rect.right, self.rect.centery))

# 坚果墙
class WallNut(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, cost=50, hp=400)

# 大嘴花
class Chomper(Plant):
    def __init__(self, x, y):
        super().__init__(x, y, cost=150, hp=100)
        self.eat_timer = 0
        self.eat_interval = 4000

    def update(self, dt, zombie_group):
        self.eat_timer += dt
        if self.eat_timer >= self.eat_interval:
            self.eat_timer = 0
            for z in zombie_group:
                if z.row_index == self.row_index and self.rect.x < z.rect.x < self.rect.x + 200:
                    z.hp -= 120
                    if z.hp <= 0:
                        z.kill()
                    break

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

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

class IceBullet(Bullet):
    def __init__(self, x, y):
        super().__init__(x, y, damage=20, speed=5)
        self.radius = 10

# ===================== 阳光类 =====================
class Sun(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.radius = 12
        self.rect = pygame.Rect(x - self.radius, y - self.radius, self.radius * 2, self.radius * 2)
        self.fall_speed = 1
        self.life_time = 10000

    def update(self, dt):
        self.life_time -= dt
        self.rect.y += self.fall_speed
        if self.life_time <= 0 or self.rect.bottom > HEIGHT:
            self.kill()

# ===================== 僵尸类 =====================
class Zombie(pygame.sprite.Sprite):
    def __init__(self, row):
        super().__init__()
        self.rect = pygame.Rect(WIDTH, row * GRID_H + 20, 60, 80)
        self.row_index = row
        self.hp = 100
        self.max_hp = 100
        self.base_speed = 0.6
        self.speed = self.base_speed
        self.slow_timer = 0
        self.attack_cd = 0
        self.attack_interval = 1000

    def slow_down(self, duration=2000):
        self.speed = self.base_speed * 0.4
        self.slow_timer = duration

    def update(self, dt, plant_group):
        if self.slow_timer > 0:
            self.slow_timer -= dt
        else:
            self.speed = self.base_speed

        hit_plant = pygame.sprite.spritecollide(self, plant_group, False)
        if hit_plant:
            self.attack_cd += dt
            if self.attack_cd >= self.attack_interval:
                self.attack_cd = 0
                for p in hit_plant:
                    p.hp -= 10
                    if p.hp <= 0:
                        p.kill()
        else:
            self.rect.x -= self.speed
        if self.rect.right <= 0:
            global game_over
            game_over = True

class ConeZombie(Zombie):
    def __init__(self, row):
        super().__init__(row)
        self.hp = 200
        self.max_hp = 200
        self.base_speed = 0.5
        self.speed = self.base_speed

class BucketZombie(Zombie):
    def __init__(self, row):
        super().__init__(row)
        self.hp = 400
        self.max_hp = 400
        self.base_speed = 0.4
        self.speed = self.base_speed

# ===================== 绘制草坪 =====================
def draw_grass():
    for r in range(ROW_NUM):
        for c in range(COL_NUM):
            x = c * GRID_W
            y = r * GRID_H
            rect = pygame.Rect(x, y, GRID_W, GRID_H)
            if (r + c) % 2 == 0:
                pygame.draw.rect(screen, GREEN, rect)
            else:
                pygame.draw.rect(screen, GREEN_LIGHT, rect)
            pygame.draw.rect(screen, BLACK, rect, 1)

# ===================== 主游戏 =====================
def main():
    global game_over
    game_state = "menu"
    game_over = False
    sun_num = 200
    selected_plant = None

    plant_group = pygame.sprite.Group()
    bullet_group = pygame.sprite.Group()
    sun_group = pygame.sprite.Group()
    zombie_group = pygame.sprite.Group()

    zombie_spawn_timer = 0
    zombie_spawn_gap = 6000   # 缩短刷新间隔，更容易看见僵尸
    difficulty_timer = 0

    while True:
        dt = clock.tick(FPS)
        screen.fill(SKY_BLUE)
        mx, my = pygame.mouse.get_pos()

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

            if event.type == pygame.MOUSEBUTTONDOWN:
                px, py = event.pos
                # 拾取阳光
                for sun in sun_group:
                    if sun.rect.collidepoint(px, py):
                        sun_num += 25
                        sun.kill()
                        break

                if game_state == "menu":
                    game_state = "run"

                if game_state == "run":
                    sun_btn = pygame.Rect(120, HEIGHT - 75, 55, 60)
                    pea_btn = pygame.Rect(185, HEIGHT - 75, 55, 60)
                    ice_btn = pygame.Rect(250, HEIGHT - 75, 55, 60)
                    nut_btn = pygame.Rect(315, HEIGHT - 75, 55, 60)
                    chomp_btn = pygame.Rect(380, HEIGHT - 75, 55, 60)

                    if sun_btn.collidepoint(px, py):
                        selected_plant = "sunflower"
                    if pea_btn.collidepoint(px, py):
                        selected_plant = "peashooter"
                    if ice_btn.collidepoint(px, py):
                        selected_plant = "icepea"
                    if nut_btn.collidepoint(px, py):
                        selected_plant = "wallnut"
                    if chomp_btn.collidepoint(px, py):
                        selected_plant = "chomper"

                    # 放置植物
                    if selected_plant is not None:
                        grid_col = px // GRID_W
                        grid_row = py // GRID_H
                        if 0 <= grid_col < COL_NUM and 0 <= grid_row < ROW_NUM:
                            plant_x = grid_col * GRID_W + 5
                            plant_y = grid_row * GRID_H + 5
                            new_plant_rect = pygame.Rect(plant_x, plant_y, GRID_W - 10, GRID_H - 10)

                            has_plant = False
                            for p in plant_group:
                                if p.rect.colliderect(new_plant_rect):
                                    has_plant = True
                                    break

                            if not has_plant:
                                if selected_plant == "sunflower" and sun_num >= 50:
                                    plant_group.add(SunFlower(plant_x, plant_y))
                                    sun_num -= 50
                                elif selected_plant == "peashooter" and sun_num >= 100:
                                    plant_group.add(PeaShooter(plant_x, plant_y))
                                    sun_num -= 100
                                elif selected_plant == "icepea" and sun_num >= 175:
                                    plant_group.add(IcePeaShooter(plant_x, plant_y))
                                    sun_num -= 175
                                elif selected_plant == "wallnut" and sun_num >= 50:
                                    plant_group.add(WallNut(plant_x, plant_y))
                                    sun_num -= 50
                                elif selected_plant == "chomper" and sun_num >= 150:
                                    plant_group.add(Chomper(plant_x, plant_y))
                                    sun_num -= 150

        if game_state == "menu":
            t1 = font_big.render("植物大战僵尸 完整版", True, BLACK)
            t2 = font_mid.render("鼠标点击屏幕任意位置开始游戏", True, RED)
            t3 = font_small.render("僵尸、豌豆射击全部修复正常", True, BLACK)
            screen.blit(t1, (WIDTH/2 - t1.get_width()/2, 180))
            screen.blit(t2, (WIDTH/2 - t2.get_width()/2, 260))
            screen.blit(t3, (WIDTH/2 - t3.get_width()/2, 320))

        elif game_state == "run":
            draw_grass()

            # 僵尸刷新逻辑
            zombie_spawn_timer += dt
            difficulty_timer += dt
            if difficulty_timer > 30000 and zombie_spawn_gap > 4000:
                zombie_spawn_gap -= 400
                difficulty_timer = 0

            if zombie_spawn_timer >= zombie_spawn_gap:
                zombie_spawn_timer = 0
                row = random.randint(0, ROW_NUM - 1)
                roll = random.random()
                if roll < 0.55:
                    new_zom = Zombie(row)
                elif roll < 0.80:
                    new_zom = ConeZombie(row)
                else:
                    new_zom = BucketZombie(row)
                zombie_group.add(new_zom)

            # 更新所有植物
            for plant in plant_group.sprites():
                if isinstance(plant, SunFlower):
                    plant.update(dt, sun_group)
                elif isinstance(plant, PeaShooter):
                    plant.update(dt, bullet_group, zombie_group)
                elif isinstance(plant, IcePeaShooter):
                    plant.update(dt, bullet_group, zombie_group)
                elif isinstance(plant, Chomper):
                    plant.update(dt, zombie_group)

            bullet_group.update()
            sun_group.update(dt)
            zombie_group.update(dt, plant_group)

            # 子弹命中僵尸，寒冰子弹附带减速
            hit_info = pygame.sprite.groupcollide(zombie_group, bullet_group, False, True)
            for zom, bullets in hit_info.items():
                for b in bullets:
                    zom.hp -= b.damage
                    if isinstance(b, IceBullet):
                        zom.slow_down()
                    if zom.hp <= 0:
                        zom.kill()

            # 绘制植物
            for p in plant_group:
                if isinstance(p, SunFlower):
                    pygame.draw.circle(screen, YELLOW, p.rect.center, 30)
                    pygame.draw.circle(screen, (255, 160, 0), p.rect.center, 15)
                elif isinstance(p, PeaShooter):
                    pygame.draw.rect(screen, DARK_GREEN, p.rect)
                    pygame.draw.circle(screen, (0, 220, 0), (p.rect.right, p.rect.centery), 12)
                elif isinstance(p, IcePeaShooter):
                    pygame.draw.rect(screen, (0, 120, 180), p.rect)
                    pygame.draw.circle(screen, ICE_BLUE, (p.rect.right, p.rect.centery), 12)
                elif isinstance(p, WallNut):
                    pygame.draw.rect(screen, (160, 110, 50), p.rect)
                    pygame.draw.circle(screen, (200, 140, 70), p.rect.center, 25)
                elif isinstance(p, Chomper):
                    pygame.draw.rect(screen, (80, 180, 80), p.rect)
                    pygame.draw.circle(screen, (255, 80, 80), (p.rect.right - 10, p.rect.centery), 18)

                hp_rate = p.hp / p.max_hp
                pygame.draw.rect(screen, RED, (p.rect.x, p.rect.y - 8, p.rect.width, 4))
                pygame.draw.rect(screen, (0, 200, 0), (p.rect.x, p.rect.y - 8, p.rect.width * hp_rate, 4))

            # 绘制子弹
            for b in bullet_group:
                color = ICE_BLUE if isinstance(b, IceBullet) else (0, 200, 255)
                pygame.draw.circle(screen, color, b.rect.center, b.radius)

            # 绘制阳光
            for s in sun_group:
                pygame.draw.circle(screen, YELLOW, s.rect.center, s.radius)

            # 绘制僵尸
            for z in zombie_group:
                if isinstance(z, BucketZombie):
                    pygame.draw.rect(screen, GRAY, z.rect)
                    pygame.draw.circle(screen, (80, 80, 80), (z.rect.centerx, z.rect.top + 15), 18)
                elif isinstance(z, ConeZombie):
                    pygame.draw.rect(screen, BROWN, z.rect)
                    pygame.draw.circle(screen, ORANGE, (z.rect.centerx, z.rect.top + 15), 16)
                else:
                    pygame.draw.rect(screen, BROWN, z.rect)

                hp_rate = z.hp / z.max_hp
                pygame.draw.rect(screen, RED, (z.rect.x, z.rect.y - 6, z.rect.width, 4))
                pygame.draw.rect(screen, (0, 180, 0), (z.rect.x, z.rect.y - 6, z.rect.width * hp_rate, 4))

            # 底部UI
            pygame.draw.rect(screen, (80, 80, 80), (0, HEIGHT - 80, WIDTH, 80))
            sun_text = font_mid.render(f"阳光：{sun_num}", True, YELLOW)
            screen.blit(sun_text, (20, HEIGHT - 70))

            sun_btn = pygame.Rect(120, HEIGHT - 75, 55, 60)
            pea_btn = pygame.Rect(185, HEIGHT - 75, 55, 60)
            ice_btn = pygame.Rect(250, HEIGHT - 75, 55, 60)
            nut_btn = pygame.Rect(315, HEIGHT - 75, 55, 60)
            chomp_btn = pygame.Rect(380, HEIGHT - 75, 55, 60)

            pygame.draw.rect(screen, YELLOW, sun_btn)
            pygame.draw.rect(screen, DARK_GREEN, pea_btn)
            pygame.draw.rect(screen, (0, 120, 180), ice_btn)
            pygame.draw.rect(screen, (160, 110, 50), nut_btn)
            pygame.draw.rect(screen, (80, 180, 80), chomp_btn)

            screen.blit(font_small.render("50", True, BLACK), (sun_btn.centerx - 8, sun_btn.bottom - 22))
            screen.blit(font_small.render("100", True, WHITE), (pea_btn.centerx - 10, pea_btn.bottom - 22))
            screen.blit(font_small.render("175", True, WHITE), (ice_btn.centerx - 12, ice_btn.bottom - 22))
            screen.blit(font_small.render("50", True, BLACK), (nut_btn.centerx - 8, nut_btn.bottom - 22))
            screen.blit(font_small.render("150", True, WHITE), (chomp_btn.centerx - 10, chomp_btn.bottom - 22))

            if selected_plant == "sunflower":
                pygame.draw.rect(screen, WHITE, sun_btn, 3)
            elif selected_plant == "peashooter":
                pygame.draw.rect(screen, WHITE, pea_btn, 3)
            elif selected_plant == "icepea":
                pygame.draw.rect(screen, WHITE, ice_btn, 3)
            elif selected_plant == "wallnut":
                pygame.draw.rect(screen, WHITE, nut_btn, 3)
            elif selected_plant == "chomper":
                pygame.draw.rect(screen, WHITE, chomp_btn, 3)

            if game_over:
                game_state = "over"

        elif game_state == "over":
            over_text = font_big.render("游戏失败！僵尸吃掉你的脑子了", True, RED)
            tip_text = font_mid.render("关闭窗口退出游戏", True, BLACK)
            screen.blit(over_text, (WIDTH/2 - over_text.get_width()/2, 220))
            screen.blit(tip_text, (WIDTH/2 - tip_text.get_width()/2, 300))

        pygame.display.flip()

if __name__ == "__main__":
    main()