import pygame
import random
import sys

# --- 初始化 ---
pygame.init()

# --- 常量定义 ---
SCREEN_WIDTH = 1000
SCREEN_HEIGHT = 600
GRID_ROWS = 5
GRID_COLS = 9
CELL_WIDTH = 80
CELL_HEIGHT = 100
GRID_OFFSET_X = 200  # 左侧留出空间放UI
GRID_OFFSET_Y = 50

FPS = 60

# 颜色定义
COLOR_BG = (30, 30, 30)
COLOR_GRID_LIGHT = (50, 200, 50)
COLOR_GRID_DARK = (40, 180, 40)
COLOR_SUN = (255, 255, 0)
COLOR_UI_BG = (20, 20, 20)
COLOR_TEXT = (255, 255, 255)

# 植物/僵尸属性
SUN_VALUE = 25
SUN_COST_SUNFLOWER = 50
SUN_COST_PEASHOOTER = 100
PEA_DAMAGE = 20
ZOMBIE_HEALTH = 100
ZOMBIE_DAMAGE = 1 # 每帧伤害

class Sun:
    def __init__(self, x, y, value=25):
        self.rect = pygame.Rect(x, y, 40, 40)
        self.value = value
        self.life = 500 # 存在时间
        self.target_y = y # 如果是天空掉落，目标位置

    def update(self):
        self.life -= 1
        # 简单的掉落动画逻辑（如果是天空掉落）
        if self.rect.y < self.target_y:
            self.rect.y += 2
        return self.life > 0

    def draw(self, surface):
        pygame.draw.circle(surface, COLOR_SUN, self.rect.center, 20)
        # 光芒
        pygame.draw.circle(surface, (255, 200, 0), self.rect.center, 15)

class Bullet:
    def __init__(self, x, y):
        self.rect = pygame.Rect(x, y, 10, 10)
        self.speed = 7
        self.damage = PEA_DAMAGE
        self.active = True

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

    def draw(self, surface):
        pygame.draw.circle(surface, (0, 255, 0), self.rect.center, 5)

class Plant:
    def __init__(self, row, col, type):
        self.row = row
        self.col = col
        self.type = type # 'sunflower' or 'peashooter'
        self.x = GRID_OFFSET_X + col * CELL_WIDTH + CELL_WIDTH // 2
        self.y = GRID_OFFSET_Y + row * CELL_HEIGHT + CELL_HEIGHT // 2
        self.rect = pygame.Rect(self.x - 30, self.y - 30, 60, 60)
        self.health = 100
        self.timer = 0

    def update(self, game):
        self.timer += 1
        
        if self.type == 'sunflower':
            # 每 5 秒产生阳光
            if self.timer % (FPS * 5) == 0:
                game.add_sun(self.rect.centerx, self.rect.centery)
        
        elif self.type == 'peashooter':
            # 检测该行是否有僵尸
            has_zombie = False
            for zombie in game.zombies:
                if zombie.row == self.row and zombie.rect.centerx > self.rect.centerx:
                    has_zombie = True
                    break
            
            if has_zombie:
                # 每 1.5 秒发射子弹
                if self.timer % (FPS * 1.5) == 0:
                    game.bullets.append(Bullet(self.rect.right, self.rect.centery))

    def draw(self, surface):
        if self.type == 'sunflower':
            # 绘制向日葵 (黄色圆形)
            pygame.draw.circle(surface, (255, 255, 0), (self.x, self.y), 25)
            pygame.draw.circle(surface, (139, 69, 19), (self.x, self.y), 15) # 花蕊
        elif self.type == 'peashooter':
            # 绘制豌豆射手 (绿色圆形+管)
            pygame.draw.circle(surface, (0, 255, 0), (self.x, self.y), 25)
            pygame.draw.rect(surface, (0, 200, 0), (self.x, self.y - 10, 40, 20)) # 炮管

class Zombie:
    def __init__(self, row):
        self.row = row
        self.x = SCREEN_WIDTH
        self.y = GRID_OFFSET_Y + row * CELL_HEIGHT + CELL_HEIGHT // 2
        self.rect = pygame.Rect(self.x - 30, self.y - 40, 60, 80)
        self.health = ZOMBIE_HEALTH
        self.speed = 1
        self.is_eating = False

    def update(self, plants):
        self.is_eating = False
        # 移动逻辑
        target_x = self.x - self.speed
        
        # 碰撞检测：是否碰到植物
        for plant in plants:
            if plant.row == self.row:
                # 简单的矩形碰撞
                if self.rect.colliderect(plant.rect):
                    self.is_eating = True
                    plant.health -= ZOMBIE_DAMAGE
                    target_x = self.x # 停下来吃
                    break
        
        if not self.is_eating:
            self.x = target_x
            self.rect.centerx = int(self.x)

    def draw(self, surface):
        # 绘制僵尸 (紫色矩形)
        pygame.draw.rect(surface, (100, 50, 100), self.rect)
        # 眼睛
        pygame.draw.circle(surface, (255, 0, 0), (self.rect.left + 15, self.rect.top + 20), 5)
        # 血条
        hp_percent = self.health / ZOMBIE_HEALTH
        pygame.draw.rect(surface, (255, 0, 0), (self.rect.x, self.rect.top - 10, 60, 5))
        pygame.draw.rect(surface, (0, 255, 0), (self.rect.x, self.rect.top - 10, 60 * hp_percent, 5))

class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
        pygame.display.set_caption("Python 植物大战僵尸")
        self.clock = pygame.time.Clock()
        
        self.sun = 150 # 初始阳光
        self.plants = []
        self.zombies = []
        self.bullets = []
        self.sun_objects = []
        
        self.zombie_spawn_timer = 0
        self.game_over = False

    def add_sun(self, x, y):
        self.sun_objects.append(Sun(x, y))

    def spawn_zombie(self):
        row = random.randint(0, GRID_ROWS - 1)
        self.zombies.append(Zombie(row))

    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
            
            if event.type == pygame.MOUSEBUTTONDOWN:
                if self.game_over:
                    # 重启游戏
                    self.__init__()
                    return True

                mx, my = pygame.mouse.get_pos()
                
                # 检查是否点击了阳光
                clicked_sun = False
                for sun in self.sun_objects[:]:
                    if sun.rect.collidepoint(mx, my):
                        self.sun += sun.value
                        self.sun_objects.remove(sun)
                        clicked_sun = True
                        break
                
                if clicked_sun: continue

                # 检查是否点击了草坪种植
                if (GRID_OFFSET_X <= mx < SCREEN_WIDTH and 
                    GRID_OFFSET_Y <= my < GRID_OFFSET_Y + GRID_ROWS * CELL_HEIGHT):
                    
                    col = (mx - GRID_OFFSET_X) // CELL_WIDTH
                    row = (my - GRID_OFFSET_Y) // CELL_HEIGHT
                    
                    # 检查该位置是否已有植物
                    occupied = False
                    for plant in self.plants:
                        if plant.row == row and plant.col == col:
                            occupied = True
                            break
                    
                    if not occupied:
                        # 尝试种植向日葵
                        if event.button == 1: # 左键
                            if self.sun >= SUN_COST_SUNFLOWER:
                                self.plants.append(Plant(row, col, 'sunflower'))
                                self.sun -= SUN_COST_SUNFLOWER
                        # 尝试种植豌豆射手 (右键或按住Shift+左键，这里简化为数字键切换)
                        # 为了简单，这里只用左键，通过UI按钮切换模式，或者直接用数字键
                        # 这里演示用数字键 1:向日葵, 2:豌豆
                        # 需要在事件循环外处理按键状态，或者增加UI按钮
                        # 简化版：左键种向日葵，右键种豌豆
                        elif event.button == 3: # 右键
                            if self.sun >= SUN_COST_PEASHOOTER:
                                self.plants.append(Plant(row, col, 'peashooter'))
                                self.sun -= SUN_COST_PEASHOOTER

            if event.type == pygame.KEYDOWN:
                 if self.game_over and event.key == pygame.K_r:
                     self.__init__()

        # 键盘连续检测（用于选择植物类型）
        keys = pygame.key.get_pressed()
        # 这里可以扩展植物选择逻辑，简化版暂略

        return True

    def update(self):
        if self.game_over: return

        # 生成僵尸
        self.zombie_spawn_timer += 1
        # 随时间增加难度
        spawn_rate = max(100, 300 - self.zombie_spawn_timer // 10)
        if self.zombie_spawn_timer % int(spawn_rate) == 0:
            self.spawn_zombie()

        # 更新植物
        for plant in self.plants:
            plant.update(self)
            if plant.health <= 0:
                self.plants.remove(plant)

        # 更新僵尸
        for zombie in self.zombies:
            zombie.update(self.plants)
            if zombie.health <= 0:
                if zombie in self.zombies:
                    self.zombies.remove(zombie)
            if zombie.rect.left < GRID_OFFSET_X:
                self.game_over = True

        # 更新子弹
        for bullet in self.bullets[:]:
            bullet.update()
            if not bullet.active:
                if bullet in self.bullets:
                    self.bullets.remove(bullet)
            else:
                # 子弹击中僵尸
                for zombie in self.zombies:
                    if zombie.row == (bullet.rect.centery - GRID_OFFSET_Y) // CELL_HEIGHT: # 简单行检测
                         if bullet.rect.colliderect(zombie.rect):
                            zombie.health -= bullet.damage
                            if bullet in self.bullets:
                                self.bullets.remove(bullet)
                            break

        # 更新阳光
        self.sun_objects = [s for s in self.sun_objects if s.update()]

    def draw(self):
        self.screen.fill(COLOR_BG)
        
        # 绘制草坪网格
        for row in range(GRID_ROWS):
            for col in range(GRID_COLS):
                x = GRID_OFFSET_X + col * CELL_WIDTH
                y = GRID_OFFSET_Y + row * CELL_HEIGHT
                color = COLOR_GRID_LIGHT if (row + col) % 2 == 0 else COLOR_GRID_DARK
                pygame.draw.rect(self.screen, color, (x, y, CELL_WIDTH, CELL_HEIGHT))
                pygame.draw.rect(self.screen, (0, 50, 0), (x, y, CELL_WIDTH, CELL_HEIGHT), 1)

        # 绘制实体
        for plant in self.plants:
            plant.draw(self.screen)
        
        for zombie in self.zombies:
            zombie.draw(self.screen)
            
        for bullet in self.bullets:
            bullet.draw(self.screen)
            
        for sun in self.sun_objects:
            sun.draw(self.screen)

        # 绘制 UI
        pygame.draw.rect(self.screen, COLOR_UI_BG, (0, 0, GRID_OFFSET_X, SCREEN_HEIGHT))
        
        # 阳光数值
        font = pygame.font.Font(None, 40)
        sun_text = font.render(f"阳光: {self.sun}", True, COLOR_SUN)
        self.screen.blit(sun_text, (20, 50))
        
        # 提示文字
        hint_font = pygame.font.Font(None, 24)
        info1 = hint_font.render("左键: 向日葵 (50)", True, COLOR_TEXT)
        info2 = hint_font.render("右键: 豌豆射手 (100)", True, COLOR_TEXT)
        self.screen.blit(info1, (20, 150))
        self.screen.blit(info2, (20, 180))

        if self.game_over:
            s = pygame.Surface((SCREEN_WIDTH, SCREEN_HEIGHT))
            s.set_alpha(200)
            s.fill((0, 0, 0))
            self.screen.blit(s, (0, 0))
            draw_text(self.screen, "僵尸吃掉了你的脑子!", 60, SCREEN_WIDTH//2, SCREEN_HEIGHT//2, (255, 0, 0))
            draw_text(self.screen, "按 [R] 或点击重新开始", 30, SCREEN_WIDTH//2, SCREEN_HEIGHT//2 + 60)

        pygame.display.flip()

def draw_text(surface, text, size, x, y, color, center=True):
    font = pygame.font.Font(None, size)
    render = font.render(text, True, color)
    rect = render.get_rect()
    if center: rect.center = (x, y)
    else: rect.topleft = (x, y)
    surface.blit(render, rect)

def main():
    game = Game()
    running = True
    while running:
        running = game.handle_events()
        game.update()
        game.draw()
        game.clock.tick(FPS)
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()