import pygame
import math
import random

# 初始化
pygame.init()
W, H = 900, 600
screen = pygame.display.set_mode((W, H))
pygame.display.set_caption("2D俯视角射击 - 多行为敌人AI")
clock = pygame.time.Clock()
FPS = 60

# 颜色
BG_COLOR = (35, 35, 35)
PLAYER_COL = (0, 210, 100)
ENEMY_COL = (220, 40, 40)
BULLET_COL = (255, 230, 0)
WHITE = (255, 255, 255)

# -------------------------- 玩家类 --------------------------
class Player:
    def __init__(self):
        self.x = W // 2
        self.y = H // 2
        self.radius = 18
        self.speed = 240
        self.hp = 100

    def update(self, dt):
        keys = pygame.key.get_pressed()
        if keys[pygame.K_w]:
            self.y -= self.speed * dt
        if keys[pygame.K_s]:
            self.y += self.speed * dt
        if keys[pygame.K_a]:
            self.x -= self.speed * dt
        if keys[pygame.K_d]:
            self.x += self.speed * dt

        # 边界锁住
        self.x = max(self.radius, min(W - self.radius, self.x))
        self.y = max(self.radius, min(H - self.radius, self.y))

    def draw(self):
        pygame.draw.circle(screen, PLAYER_COL, (int(self.x), int(self.y)), self.radius)
        # 瞄准线
        mx, my = pygame.mouse.get_pos()
        pygame.draw.line(screen, WHITE, (self.x, self.y), (mx, my), 2)

# -------------------------- 子弹类 --------------------------
class Bullet:
    def __init__(self, x, y, angle):
        self.x = x
        self.y = y
        self.angle = angle
        self.speed = 580
        self.radius = 4

    def update(self, dt):
        self.x += math.cos(self.angle) * self.speed * dt
        self.y += math.sin(self.angle) * self.speed * dt

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

    def out_range(self):
        return self.x < 0 or self.x > W or self.y < 0 or self.y > H

# -------------------------- 敌人类（5种行为AI） --------------------------
class Enemy:
    def __init__(self):
        # 四边随机出生
        side = random.randint(0, 3)
        if side == 0:
            self.x = random.randint(0, W)
            self.y = -40
        elif side == 1:
            self.x = random.randint(0, W)
            self.y = H + 40
        elif side == 2:
            self.x = -40
            self.y = random.randint(0, H)
        else:
            self.x = W + 40
            self.y = random.randint(0, H)

        self.radius = 19
        self.speed = 75
        self.hp = 2

        # 随机分配行为类型
        self.behavior = random.choice(["normal", "patrol", "ambush", "rush", "tank"])
        self.walk_timer = 0
        self.walk_dir = random.uniform(0, math.pi * 2)

    def update(self, dt, px, py):
        dx = px - self.x
        dy = py - self.y
        dist = math.hypot(dx, dy)

        # 1. 普通追击
        if self.behavior == "normal":
            ang = math.atan2(dy, dx)
            self.x += math.cos(ang) * self.speed * dt
            self.y += math.sin(ang) * self.speed * dt

        # 2. 巡逻 + 进入视野才追击
        elif self.behavior == "patrol":
            if dist < 260:
                ang = math.atan2(dy, dx)
                self.x += math.cos(ang) * self.speed * dt
                self.y += math.sin(ang) * self.speed * dt
            else:
                self.walk_timer += dt
                if self.walk_timer > 0.9:
                    self.walk_dir = random.uniform(0, math.pi * 2)
                    self.walk_timer = 0
                self.x += math.cos(self.walk_dir) * 35 * dt
                self.y += math.sin(self.walk_dir) * 35 * dt

        # 3. 伏击型：靠近才冲
        elif self.behavior == "ambush":
            if dist < 130:
                ang = math.atan2(dy, dx)
                self.x += math.cos(ang) * self.speed * 1.35 * dt
                self.y += math.sin(ang) * self.speed * 1.35 * dt

        # 4. 快速突袭：距离越近越快 + 小幅飘忽
        elif self.behavior == "rush":
            ang = math.atan2(dy, dx)
            ang += random.uniform(-0.13, 0.13)
            spd_mod = 0.55 + dist / 320
            self.x += math.cos(ang) * self.speed * spd_mod * dt
            self.y += math.sin(ang) * self.speed * spd_mod * dt

        # 5. 肉盾慢速
        elif self.behavior == "tank":
            ang = math.atan2(dy, dx)
            self.x += math.cos(ang) * self.speed * 0.5 * dt
            self.y += math.sin(ang) * self.speed * 0.5 * dt

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

# -------------------------- 游戏初始化 --------------------------
player = Player()
bullets = []
enemies = []
spawn_timer = 0
score = 0
font = pygame.font.SysFont(None, 36)

# 主循环
running = True
while running:
    dt = clock.tick(FPS) / 1000.0
    now = pygame.time.get_ticks()
    screen.fill(BG_COLOR)

    # 事件
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        # 鼠标左键射击
        if event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
            mx, my = pygame.mouse.get_pos()
            ang = math.atan2(my - player.y, mx - player.x)
            bullets.append(Bullet(player.x, player.y, ang))

    # 更新玩家
    player.update(dt)

    # 生成敌人
    if now - spawn_timer > 1400:
        spawn_timer = now
        enemies.append(Enemy())

    # 更新子弹
    for b in bullets[:]:
        b.update(dt)
        if b.out_range():
            bullets.remove(b)

    # 更新敌人 & 碰撞逻辑
    for e in enemies[:]:
        e.update(dt, player.x, player.y)

        # 子弹击中敌人
        for b in bullets[:]:
            d = math.hypot(b.x - e.x, b.y - e.y)
            if d < e.radius + b.radius:
                e.hp -= 1
                bullets.remove(b)
                if e.hp <= 0:
                    enemies.remove(e)
                    score += 15
                break

        # 敌人碰到玩家扣血
        d_player = math.hypot(e.x - player.x, e.y - player.y)
        if d_player < e.radius + player.radius:
            player.hp -= 0.4

    # 绘制
    player.draw()
    for b in bullets:
        b.draw()
    for e in enemies:
        e.draw()

    # UI文字
    screen.blit(font.render(f"血量: {int(player.hp)}", True, WHITE), (20, 20))
    screen.blit(font.render(f"分数: {score}", True, WHITE), (20, 60))

    # 游戏结束
    if player.hp <= 0:
        over_text = font.render("游戏结束!", True, (255, 0, 0))
        screen.blit(over_text, (W//2 - 70, H//2))

    pygame.display.flip()

pygame.quit()