import pygame
import sys
import math
import random

# --- 初始化 ---
pygame.init()
WIDTH, HEIGHT = 600, 800
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Pygame 飞行躲避 - Arcade Shooter")
clock = pygame.time.Clock()

# 颜色
BG_COLOR = (20, 20, 40)
PLAYER_COLOR = (100, 200, 255)
ENEMY_COLOR = (255, 80, 80)
BULLET_COLOR = (255, 255, 100)
ENEMY_BULLET_COLOR = (255, 100, 100)
TEXT_COLOR = (255, 255, 255)
FLAME_COLOR = (255, 150, 50)

FONT = pygame.font.SysFont("consolas", 30)
BIG_FONT = pygame.font.SysFont("consolas", 60)

# 配置
PLAYER_SIZE = 20
BULLET_SPEED = 10
ENEMY_BULLET_SPEED = 4
PLAYER_SPEED = 6
FRICTION = 0.92  # 惯性摩擦

# --- 游戏对象 ---

class Player:
    def __init__(self):
        self.x = WIDTH // 2
        self.y = HEIGHT - 100
        self.vx = 0
        self.vy = 0
        self.shoot_cd = 0
        self.hp = 3
        self.flash_timer = 0

    def update(self, keys, dt):
        # 输入
        ax, ay = 0, 0
        if keys[pygame.K_w]: ay -= 1
        if keys[pygame.K_s]: ay += 1
        if keys[pygame.K_a]: ax -= 1
        if keys[pygame.K_d]: ax += 1

        # 加速度
        self.vx += ax * 0.8
        self.vy += ay * 0.8

        # 摩擦
        self.vx *= FRICTION
        self.vy *= FRICTION

        # 限速
        speed = math.hypot(self.vx, self.vy)
        if speed > PLAYER_SPEED:
            self.vx = (self.vx / speed) * PLAYER_SPEED
            self.vy = (self.vy / speed) * PLAYER_SPEED

        self.x += self.vx
        self.y += self.vy

        # 边界
        self.x = max(PLAYER_SIZE, min(WIDTH - PLAYER_SIZE, self.x))
        self.y = max(PLAYER_SIZE, min(HEIGHT - PLAYER_SIZE, self.y))

        if self.shoot_cd > 0: self.shoot_cd -= 1
        if self.flash_timer > 0: self.flash_timer -= 1

    def shoot(self):
        if self.shoot_cd <= 0:
            self.shoot_cd = 8
            return Bullet(self.x, self.y - PLAYER_SIZE, 0, -BULLET_SPEED)
        return None

    def hit(self):
        if self.flash_timer > 0: return False
        self.hp -= 1
        self.flash_timer = 60
        return True

    def draw(self, surface):
        # 尾焰
        flame_len = 10 + abs(self.vy) * 3
        pygame.draw.polygon(surface, FLAME_COLOR, [
            (self.x - 8, self.y + PLAYER_SIZE),
            (self.x + 8, self.y + PLAYER_SIZE),
            (self.x, self.y + PLAYER_SIZE + flame_len)
        ])

        # 机身
        color = (255, 100, 100) if self.flash_timer > 0 and self.flash_timer % 10 < 5 else PLAYER_COLOR
        pygame.draw.polygon(surface, color, [
            (self.x, self.y - PLAYER_SIZE),
            (self.x - PLAYER_SIZE, self.y + PLAYER_SIZE),
            (self.x + PLAYER_SIZE, self.y + PLAYER_SIZE)
        ])

class Bullet:
    def __init__(self, x, y, vx, vy, is_enemy=False):
        self.x = x
        self.y = y
        self.vx = vx
        self.vy = vy
        self.is_enemy = is_enemy
        self.alive = True
        self.trail = []

    def update(self):
        self.x += self.vx
        self.y += self.vy
        self.trail.append((self.x, self.y))
        if len(self.trail) > 5: self.trail.pop(0)

        if self.x < 0 or self.x > WIDTH or self.y < 0 or self.y > HEIGHT:
            self.alive = False

    def draw(self, surface):
        # 拖尾
        for i, (tx, ty) in enumerate(self.trail):
            alpha = int(150 * (i / len(self.trail)))
            r = int(4 * (i / len(self.trail)))
            color = ENEMY_BULLET_COLOR if self.is_enemy else BULLET_COLOR
            pygame.draw.circle(surface, color, (int(tx), int(ty)), max(2, r))

        color = ENEMY_BULLET_COLOR if self.is_enemy else BULLET_COLOR
        pygame.draw.circle(surface, color, (int(self.x), int(self.y)), 5)

class Enemy:
    def __init__(self, x, y, pattern=0):
        self.x = x
        self.y = y
        self.hp = 2
        self.pattern = pattern  # 0=直线, 1=扇形, 2=追踪
        self.shoot_cd = random.randint(30, 90)
        self.vx = random.uniform(-1, 1)
        self.vy = 1.5
        self.alive = True
        self.flash_timer = 0

    def update(self, player):
        self.x += self.vx
        self.y += self.vy

        # 左右边界反弹
        if self.x < 30 or self.x > WIDTH - 30:
            self.vx *= -1

        # 出界销毁
        if self.y > HEIGHT + 50:
            self.alive = False

        if self.shoot_cd > 0: self.shoot_cd -= 1
        if self.flash_timer > 0: self.flash_timer -= 1

    def shoot(self, player):
        if self.shoot_cd <= 0:
            self.shoot_cd = random.randint(40, 100)
            bullets = []

            if self.pattern == 0:  # 直线
                bullets.append(Bullet(self.x, self.y + 20, 0, ENEMY_BULLET_SPEED, True))
            elif self.pattern == 1:  # 扇形
                for angle in [-0.3, 0, 0.3]:
                    vx = math.sin(angle) * ENEMY_BULLET_SPEED
                    vy = math.cos(angle) * ENEMY_BULLET_SPEED
                    bullets.append(Bullet(self.x, self.y + 20, vx, vy, True))
            elif self.pattern == 2:  # 追踪
                dx = player.x - self.x
                dy = player.y - self.y
                dist = math.hypot(dx, dy)
                if dist > 0:
                    vx = (dx / dist) * ENEMY_BULLET_SPEED * 0.8
                    vy = (dy / dist) * ENEMY_BULLET_SPEED * 0.8
                    bullets.append(Bullet(self.x, self.y + 20, vx, vy, True))

            return bullets
        return []

    def hit(self):
        self.hp -= 1
        self.flash_timer = 6
        if self.hp <= 0:
            self.alive = False
            return True
        return False

    def draw(self, surface):
        color = (255, 200, 200) if self.flash_timer > 0 else ENEMY_COLOR
        pygame.draw.polygon(surface, color, [
            (self.x, self.y + 20),
            (self.x - 20, self.y - 15),
            (self.x + 20, self.y - 15)
        ])

# --- 主程序 ---

def main():
    player = Player()
    bullets = []
    enemies = []
    score = 0
    frame = 0
    game_over = False
    shake_timer = 0

    running = True
    while running:
        clock.tick(60)
        frame += 1
        dt = 1

        for event in pygame.event.get():
            if event.type == pygame.QUIT: pygame.quit(); sys.exit()
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE: pygame.quit(); sys.exit()
                if event.key == pygame.K_r and game_over:
                    player = Player()
                    bullets = []
                    enemies = []
                    score = 0
                    frame = 0
                    game_over = False

        if not game_over:
            keys = pygame.key.get_pressed()
            player.update(keys, dt)

            # 自动射击
            if keys[pygame.K_SPACE]:
                b = player.shoot()
                if b: bullets.append(b)

            # 生成敌人
            if frame % 60 == 0:
                x = random.randint(50, WIDTH - 50)
                pattern = random.choice([0, 0, 1, 1, 2])  # 权重
                enemies.append(Enemy(x, -30, pattern))

            # 更新敌人
            for e in enemies[:]:
                e.update(player)
                new_bullets = e.shoot(player)
                bullets.extend(new_bullets)
                if not e.alive:
                    enemies.remove(e)

            # 更新子弹
            for b in bullets[:]:
                b.update()
                if not b.alive:
                    bullets.remove(b)
                    continue

                # 玩家子弹打敌人
                if not b.is_enemy:
                    for e in enemies:
                        if e.alive and math.hypot(b.x - e.x, b.y - e.y) < 25:
                            if e.hit():
                                score += 10
                            b.alive = False
                            break

                # 敌人子弹打玩家
                if b.is_enemy:
                    if math.hypot(b.x - player.x, b.y - player.y) < PLAYER_SIZE:
                        if player.hit():
                            shake_timer = 10
                            b.alive = False
                            if player.hp <= 0:
                                game_over = True

            # 敌人撞玩家
            for e in enemies:
                if e.alive and math.hypot(e.x - player.x, e.y - player.y) < PLAYER_SIZE + 20:
                    if player.hit():
                        shake_timer = 10
                        e.hit()
                        if player.hp <= 0:
                            game_over = True

        if shake_timer > 0: shake_timer -= 1

        # --- 绘图 ---
        # 屏幕震动
        offset_x = random.randint(-3, 3) if shake_timer > 0 else 0
        offset_y = random.randint(-3, 3) if shake_timer > 0 else 0

        screen.fill(BG_COLOR)

        # 背景星星
        for i in range(20):
            x = (i * 137 + frame * 0.5) % WIDTH
            y = (i * 251 + frame * 1.5) % HEIGHT
            pygame.draw.circle(screen, (100, 100, 150), (int(x), int(y)), 1)

        # 游戏对象（带震动偏移）
        for b in bullets:
            b.x += offset_x
            b.y += offset_y
            b.draw(screen)
            b.x -= offset_x
            b.y -= offset_y

        for e in enemies:
            e.x += offset_x
            e.y += offset_y
            e.draw(screen)
            e.x -= offset_x
            e.y -= offset_y

        player.x += offset_x
        player.y += offset_y
        player.draw(screen)
        player.x -= offset_x
        player.y -= offset_y

        # UI
        hp_txt = FONT.render(f"HP: {'♥' * player.hp}", True, (255, 100, 100))
        score_txt = FONT.render(f"Score: {score}", True, TEXT_COLOR)
        screen.blit(hp_txt, (20, 20))
        screen.blit(score_txt, (WIDTH - 150, 20))

        hint = FONT.render("WASD: Move | SPACE: Shoot | R: Restart", True, (150, 150, 150))
        screen.blit(hint, (WIDTH//2 - 200, HEIGHT - 40))

        if game_over:
            overlay = pygame.Surface((WIDTH, HEIGHT), pygame.SRCALPHA)
            overlay.fill((0, 0, 0, 150))
            screen.blit(overlay, (0, 0))
            txt = BIG_FONT.render("GAME OVER", True, (255, 80, 80))
            sub = FONT.render(f"Score: {score} | Press R to Restart", True, TEXT_COLOR)
            screen.blit(txt, txt.get_rect(center=(WIDTH//2, HEIGHT//2 - 40)))
            screen.blit(sub, sub.get_rect(center=(WIDTH//2, HEIGHT//2 + 30)))

        pygame.display.flip()

    pygame.quit()

if __name__ == "__main__":
    main()