import pygame
import random
import sys

# 初始化Pygame
pygame.init()

# 屏幕设置
SCREEN_WIDTH = 800
SCREEN_HEIGHT = 600
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption("植物大战僵尸 - 无敌全屏+无限阳光")

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

# 帧率
clock = pygame.time.Clock()
FPS = 60

# 网格道路设置
GRID_X = 150
GRID_Y = 100
CELL_SIZE = 80
ROWS = 5
COLS = 9

# 每行道路Y坐标
ROW_Y_LIST = [
    GRID_Y + 0 * CELL_SIZE + CELL_SIZE // 2,
    GRID_Y + 1 * CELL_SIZE + CELL_SIZE // 2,
    GRID_Y + 2 * CELL_SIZE + CELL_SIZE // 2,
    GRID_Y + 3 * CELL_SIZE + CELL_SIZE // 2,
    GRID_Y + 4 * CELL_SIZE + CELL_SIZE // 2
]

# 植物类（无敌 + 全屏攻击）


class Plant(pygame.sprite.Sprite):
    def __init__(self, x, y, plant_type, row):
        super().__init__()
        self.type = plant_type
        self.x = x
        self.y = y
        self.row = row

        if self.type == "sunflower":
            self.image = pygame.Surface((CELL_SIZE-10, CELL_SIZE-10))
            self.image.fill(YELLOW)
            self.cost = 0
        else:
            self.image = pygame.Surface((CELL_SIZE-10, CELL_SIZE-10))
            self.image.fill(GREEN)
            self.cost = 0

        self.rect = self.image.get_rect(center=(x, y))
        self.shoot_timer = 0

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

# 子弹类


class Bullet(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((12, 12))
        self.image.fill(BLUE)
        self.rect = self.image.get_rect(center=(x, y))
        self.speed = 12

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

# 僵尸类【固定道路行走，不跨行】


class Zombie(pygame.sprite.Sprite):
    def __init__(self):
        super().__init__()
        self.image = pygame.Surface((CELL_SIZE-10, CELL_SIZE-10))
        self.image.fill(RED)
        self.rect = self.image.get_rect()
        self.row = random.randint(0, ROWS - 1)
        self.rect.x = SCREEN_WIDTH
        self.rect.centery = ROW_Y_LIST[self.row]
        self.speed = 1
        self.hp = 10

    def update(self):
        self.rect.x -= self.speed
        self.rect.centery = ROW_Y_LIST[self.row]
        if self.rect.right < 0:
            self.kill()

# 阳光


class Sun(pygame.sprite.Sprite):
    def __init__(self, x, y):
        super().__init__()
        self.image = pygame.Surface((22, 22))
        self.image.fill(YELLOW)
        self.rect = self.image.get_rect(center=(x, y))

    def update(self):
        self.rect.y += 1
        if self.rect.top > SCREEN_HEIGHT:
            self.kill()


# 精灵组
all_sprites = pygame.sprite.Group()
plants = pygame.sprite.Group()
zombies = pygame.sprite.Group()
bullets = pygame.sprite.Group()
suns = pygame.sprite.Group()

# ========== 无限阳光 ==========
sun_count = 999999
selected_plant = "peashooter"
zombie_spawn_timer = 0

# 生成阳光


def spawn_sun(x, y):
    sun = Sun(x, y)
    all_sprites.add(sun)
    suns.add(sun)


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

    # 强制保持无限阳光
    sun_count = 999999

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

        if event.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()

            # 兼容旧版pygame 收集阳光
            clicked_sun = None
            for s in suns:
                if s.rect.collidepoint(mx, my):
                    clicked_sun = s
                    break
            if clicked_sun:
                clicked_sun.kill()

            # 免费种植物
            else:
                col = (mx - GRID_X) // CELL_SIZE
                row = (my - GRID_Y) // CELL_SIZE
                if 0 <= col < COLS and 0 <= row < ROWS:
                    x = GRID_X + col * CELL_SIZE + CELL_SIZE//2
                    y = GRID_Y + row * CELL_SIZE + CELL_SIZE//2

                    can_plant = True
                    for p in plants:
                        if abs(p.x - x) < 20 and abs(p.y - y) < 20:
                            can_plant = False
                            break

                    if can_plant:
                        plant = Plant(x, y, selected_plant, row)
                        all_sprites.add(plant)
                        plants.add(plant)

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_1:
                selected_plant = "peashooter"
            if event.key == pygame.K_2:
                selected_plant = "sunflower"

    # 生成僵尸
    zombie_spawn_timer += 1
    if zombie_spawn_timer >= 100:
        zombie = Zombie()
        all_sprites.add(zombie)
        zombies.add(zombie)
        zombie_spawn_timer = 0

    # 全屏子弹打僵尸
    hits = pygame.sprite.groupcollide(zombies, bullets, False, True)
    for zombie in hits:
        zombie.hp -= 1
        if zombie.hp <= 0:
            zombie.kill()

    # 植物无敌
    pygame.sprite.groupcollide(plants, zombies, False, False)

    # 绘制
    screen.fill(GRASS)
    for row in range(ROWS):
        for col in range(COLS):
            x = GRID_X + col * CELL_SIZE
            y = GRID_Y + row * CELL_SIZE
            pygame.draw.rect(screen, BROWN, (x, y, CELL_SIZE, CELL_SIZE), 2)

    all_sprites.update()
    all_sprites.draw(screen)

    # UI
    font = pygame.font.SysFont(None, 32)
    pygame.draw.rect(screen, BLUE, (50, 20, 100, 40))
    pygame.draw.rect(screen, YELLOW, (200, 20, 100, 40))
    screen.blit(font.render("豌豆(1)", True, WHITE), (60, 30))
    screen.blit(font.render("向日葵(2)", True, BLACK), (210, 30))
    screen.blit(font.render("阳光: 无限", True, YELLOW), (350, 30))
    screen.blit(font.render(f"已选: {selected_plant}", True, WHITE), (500, 30))

    pygame.display.flip()

pygame.quit()
sys.exit()