"""通过真实 HTTP、MySQL 和 Redis 测量注册并发性能并清理数据"""

import argparse
import asyncio
import json
import math
import os
import secrets
import socket
import subprocess
import sys
import time
from collections import Counter
from pathlib import Path

import httpx
import psutil

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))


def summarize(samples: list[float]) -> dict:
    """汇总毫秒延迟并保留尾部耗时"""
    values = sorted(samples)
    if not values:
        return {}
    return {
        "count": len(values),
        "mean_ms": round(sum(values) / len(values), 2),
        "p50_ms": round(values[math.ceil(len(values) * 0.50) - 1], 2),
        "p95_ms": round(values[math.ceil(len(values) * 0.95) - 1], 2),
        "p99_ms": round(values[math.ceil(len(values) * 0.99) - 1], 2),
        "max_ms": round(values[-1], 2),
    }


async def cleanup(prefix: str) -> dict:
    """删除本轮用户及其 Redis 会话并核验残留"""
    from sqlalchemy import delete, func, select

    from app.core.config import get_settings
    from app.core.database import async_session_factory, close_database
    from app.core.redis import close_redis, redis_client
    from app.model.user import User

    try:
        async with async_session_factory() as session:
            condition = User.username.startswith(prefix, autoescape=True)
            ids = set(map(str, (await session.scalars(select(User.id).where(condition))).all()))
            keys_deleted = 0
            pattern = f"{get_settings().redis_key_prefix}:session:*"
            async for key in redis_client.scan_iter(match=pattern, count=200):
                if await redis_client.get(key) in ids:
                    keys_deleted += await redis_client.delete(key)
            result = await session.execute(delete(User).where(condition))
            await session.commit()
            remaining = await session.scalar(
                select(func.count()).select_from(User).where(condition)
            )
            remaining_sessions = 0
            async for key in redis_client.scan_iter(match=pattern, count=200):
                if await redis_client.get(key) in ids:
                    remaining_sessions += 1
            return {
                "users_deleted": result.rowcount,
                "sessions_deleted": keys_deleted,
                "users_remaining": remaining,
                "sessions_remaining": remaining_sessions,
            }
    finally:
        await close_redis()
        await close_database()


async def baseline() -> dict:
    """分别测量热连接往返和本机密码计算作为诊断基线"""
    from sqlalchemy import text

    from app.core.database import engine
    from app.core.redis import redis_client
    from app.core.security import hash_password

    results = {}
    async with engine.connect() as connection:
        await connection.execute(text("SELECT 1"))
        samples = []
        for _ in range(10):
            start = time.perf_counter()
            await connection.execute(text("SELECT 1"))
            samples.append((time.perf_counter() - start) * 1000)
        results["mysql_select_1"] = summarize(samples)
    await redis_client.ping()
    samples = []
    for _ in range(10):
        start = time.perf_counter()
        await redis_client.ping()
        samples.append((time.perf_counter() - start) * 1000)
    results["redis_ping"] = summarize(samples)
    samples = []
    for _ in range(10):
        start = time.perf_counter()
        await hash_password("BenchmarkOnly123")
        samples.append((time.perf_counter() - start) * 1000)
    results["local_password_hash"] = summarize(samples)
    return results


async def stage(
    client: httpx.AsyncClient,
    process: psutil.Process,
    prefix: str,
    password: str,
    token: str,
    concurrency: int,
    seconds: float,
) -> dict:
    """并行注册并采样已登录请求延迟与服务资源"""
    latencies: list[float] = []
    probe_latencies: list[float] = []
    errors: Counter = Counter()
    probe_errors: Counter = Counter()
    cpu: list[float] = []
    memory: list[float] = []
    sequence = 0
    completed = 0
    stop = asyncio.Event()
    start = time.perf_counter()
    deadline = start + seconds
    processes = [process, *process.children(recursive=True)]
    for item in processes:
        item.cpu_percent()

    async def worker() -> None:
        """保持一个在途注册请求直到本阶段停止发起请求"""
        nonlocal sequence, completed
        while time.perf_counter() < deadline:
            sequence += 1
            name = f"{prefix}{concurrency}_{sequence}"
            tick = time.perf_counter()
            try:
                response = await client.post(
                    "/api/auth/register",
                    json={
                        "username": name,
                        "email": f"{name}@example.com",
                        "password": password,
                    },
                )
                if response.status_code == 200 and response.json().get("code") == 0:
                    completed += 1
                else:
                    errors[str(response.status_code)] += 1
            except httpx.HTTPError as exc:
                errors[type(exc).__name__] += 1
            latencies.append((time.perf_counter() - tick) * 1000)

    async def probe() -> None:
        """连续探测涉及数据库和 Redis 的普通认证请求"""
        while not stop.is_set():
            tick = time.perf_counter()
            try:
                response = await client.get(
                    "/api/users/me", headers={"Authorization": f"Bearer {token}"}
                )
                if response.status_code != 200 or response.json().get("code") != 0:
                    probe_errors[str(response.status_code)] += 1
            except httpx.HTTPError as exc:
                probe_errors[type(exc).__name__] += 1
            probe_latencies.append((time.perf_counter() - tick) * 1000)
            cpu.append(sum(item.cpu_percent() for item in processes if item.is_running()))
            memory.append(
                sum(item.memory_info().rss for item in processes if item.is_running()) / 1024**2
            )
            await asyncio.sleep(0.1)

    probe_task = asyncio.create_task(probe())
    try:
        await asyncio.gather(*(worker() for _ in range(concurrency)))
    finally:
        elapsed = time.perf_counter() - start
        stop.set()
        await probe_task
    return {
        "concurrency": concurrency,
        "elapsed_s": round(elapsed, 2),
        "success": completed,
        "errors": dict(errors),
        "success_rps": round(completed / elapsed, 2),
        "registration": summarize(latencies),
        "authenticated_probe": summarize(probe_latencies),
        "probe_errors": dict(probe_errors),
        "cpu_mean_percent_one_core": round(sum(cpu) / len(cpu), 1),
        "cpu_peak_percent_one_core": round(max(cpu), 1),
        "rss_peak_mb": round(max(memory), 1),
    }


async def run(args: argparse.Namespace) -> None:
    """启动独立单进程服务完成压测并确保结束后清理"""
    prefix = f"bench{secrets.token_hex(6)}"
    password = secrets.token_urlsafe(20) + "a9"
    with socket.socket() as bound:
        bound.bind(("127.0.0.1", 0))
        port = bound.getsockname()[1]
    report = {
        "run_prefix": prefix,
        "workers": 1,
        "logical_cpus": os.cpu_count(),
        "physical_cpus": psutil.cpu_count(logical=False),
        "stages": [],
    }
    process = await asyncio.to_thread(
        subprocess.Popen,
        [
            sys.executable,
            "-m",
            "uvicorn",
            "app.main:app",
            "--host",
            "127.0.0.1",
            "--port",
            str(port),
            "--no-access-log",
            "--log-level",
            "error",
        ],
        cwd=ROOT,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
        creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0),
    )
    try:
        async with httpx.AsyncClient(
            base_url=f"http://127.0.0.1:{port}",
            timeout=60,
            limits=httpx.Limits(max_connections=256),
            trust_env=False,
        ) as client:
            for _ in range(150):
                if process.poll() is not None:
                    raise RuntimeError("压测服务启动失败，请检查依赖和配置")
                try:
                    if (await client.get("/openapi.json")).status_code == 200:
                        break
                except httpx.HTTPError:
                    pass
                await asyncio.sleep(0.2)
            else:
                raise RuntimeError("压测服务启动超时")
            report["baseline"] = await baseline()
            print(json.dumps(report["baseline"], ensure_ascii=False), flush=True)
            name = prefix + "probe"
            response = await client.post(
                "/api/auth/register",
                json={
                    "username": name,
                    "email": f"{name}@example.com",
                    "password": password,
                },
            )
            if response.status_code != 200:
                raise RuntimeError(f"探针用户注册失败 HTTP {response.status_code}")
            response = await client.post(
                "/api/auth/login", json={"account": name, "password": password}
            )
            if response.status_code != 200:
                raise RuntimeError(f"探针登录失败 HTTP {response.status_code}")
            token = response.json()["data"]["token"]
            samples = []
            for _ in range(10):
                tick = time.perf_counter()
                response = await client.get(
                    "/api/users/me", headers={"Authorization": f"Bearer {token}"}
                )
                if response.status_code != 200:
                    raise RuntimeError("认证基线请求失败")
                samples.append((time.perf_counter() - tick) * 1000)
            report["baseline"]["authenticated_request"] = summarize(samples)
            for concurrency in args.concurrency:
                result = await stage(
                    client,
                    psutil.Process(process.pid),
                    prefix,
                    password,
                    token,
                    concurrency,
                    args.seconds,
                )
                report["stages"].append(result)
                print(json.dumps(result, ensure_ascii=False), flush=True)
    except Exception as exc:
        report["error_type"] = type(exc).__name__
        raise
    finally:
        children = []
        if process.poll() is None:
            children = psutil.Process(process.pid).children(recursive=True)
        for child in children:
            child.terminate()
        process.terminate()
        await asyncio.to_thread(process.wait)
        await asyncio.to_thread(psutil.wait_procs, children, timeout=10)
        try:
            report["cleanup"] = await cleanup(prefix)
        except Exception as exc:
            report["cleanup"] = {"error_type": type(exc).__name__, "verified": False}
            raise
        finally:
            args.output.parent.mkdir(parents=True, exist_ok=True)
            args.output.write_text(
                json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8"
            )
        print(json.dumps(report["cleanup"], ensure_ascii=False), flush=True)


def main() -> None:
    """解析压测档位和时长并运行真实依赖测试"""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--concurrency", nargs="+", type=int, default=[1, 4, 8, 16, 32, 64])
    parser.add_argument("--seconds", type=float, default=12)
    parser.add_argument("--output", type=Path, default=ROOT / "artifacts/load-test.json")
    args = parser.parse_args()
    if args.seconds <= 0 or any(value < 1 or value > 200 for value in args.concurrency):
        parser.error("时长必须为正数，并发必须在 1 到 200 之间")
    try:
        asyncio.run(run(args))
    except Exception as exc:
        print(f"压测失败：{type(exc).__name__}，请检查报告和服务配置", file=sys.stderr)
        raise SystemExit(1) from None


if __name__ == "__main__":
    main()