#!/usr/bin/env python3
"""MeshSafe 并发压测脚本。

用法:
    python tests/benchmark.py                              # 跑全部默认场景
    python tests/benchmark.py --concurrent 50              # 只跑 50 并发
    python tests/benchmark.py --concurrent 50 100          # 跑 50 和 100 两档
    python tests/benchmark.py --batch-sizes 10 50 100      # 批量档
    python tests/benchmark.py --base-url http://127.0.0.1:44380
    python tests/benchmark.py --report out.json            # 额外输出 JSON 报告

原理:
    httpx.AsyncClient + asyncio.Semaphore 控制并发数
    样本从 tests/samples/ 自动扫描，按大小分档
    大文件（5MB）现场 os.urandom 生成后自动清理
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
import statistics
import sys
import tempfile
import time
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
from typing import List, Optional, Tuple

import httpx

# ---------------------------------------------------------------------------
# 配置常量
# ---------------------------------------------------------------------------

DEFAULT_BASE_URL = "http://127.0.0.1:44380"
SAMPLES_DIR = Path(__file__).resolve().parent / "samples"
DEFAULT_TIMEOUT = 60.0  # 单个请求 60s 超时
WARMUP_COUNT = 5        # 预热次数（不计入统计）

# 大小分档（字节）
SIZE_TINY = 2 * 1024        # ≤ 2KB
SIZE_SMALL = 16 * 1024      # 2KB ~ 16KB
SIZE_MEDIUM = 100 * 1024    # 16KB ~ 100KB
SIZE_LARGE = 5 * 1024 * 1024  # ≥ 5MB（现场生成）


# ---------------------------------------------------------------------------
# 数据类
# ---------------------------------------------------------------------------

@dataclass
class SampleFile:
    path: Path
    name: str
    size: int
    size_bucket: str  # tiny / small / medium / large


@dataclass
class RequestResult:
    idx: int
    start_ts: float
    end_ts: float
    duration_ms: int
    status_code: int
    ok: bool
    error: Optional[str] = None
    files_count: int = 0
    avg_file_ms: int = 0


@dataclass
class ScenarioResult:
    name: str
    concurrent: int
    files_per_request: int
    size_bucket: str
    started_at: str
    finished_at: str
    total_duration_ms: int
    results: List[RequestResult] = field(default_factory=list)

    # 聚合统计
    @property
    def total_requests(self) -> int:
        return len(self.results)

    @property
    def total_files(self) -> int:
        return sum(r.files_count for r in self.results)

    @property
    def qps(self) -> float:
        if self.total_duration_ms <= 0:
            return 0.0
        return self.total_requests / (self.total_duration_ms / 1000.0)

    @property
    def files_per_sec(self) -> float:
        if self.total_duration_ms <= 0:
            return 0.0
        return self.total_files / (self.total_duration_ms / 1000.0)

    @property
    def latencies_ms(self) -> List[int]:
        return [r.duration_ms for r in self.results if r.ok]

    @property
    def avg_ms(self) -> float:
        data = self.latencies_ms
        return statistics.mean(data) if data else 0.0

    @property
    def p50_ms(self) -> float:
        data = self.latencies_ms
        return self._percentile(data, 50) if data else 0.0

    @property
    def p95_ms(self) -> float:
        data = self.latencies_ms
        return self._percentile(data, 95) if data else 0.0

    @property
    def p99_ms(self) -> float:
        data = self.latencies_ms
        return self._percentile(data, 99) if data else 0.0

    @property
    def error_rate(self) -> float:
        if not self.results:
            return 0.0
        bad = sum(1 for r in self.results if not r.ok)
        return bad / len(self.results)

    @property
    def slowest(self) -> Optional[RequestResult]:
        oks = [r for r in self.results if r.ok]
        return max(oks, key=lambda r: r.duration_ms) if oks else None

    @staticmethod
    def _percentile(data: List[int], pct: float) -> float:
        if not data:
            return 0.0
        s = sorted(data)
        k = (len(s) - 1) * (pct / 100.0)
        f = int(k)
        c = f + 1 if f + 1 < len(s) else f
        if f == c:
            return float(s[f])
        return float(s[f] * (c - k) + s[c] * (k - f))


# ---------------------------------------------------------------------------
# 样本加载
# ---------------------------------------------------------------------------

def load_samples() -> List[SampleFile]:
    """扫描 tests/samples/ 下所有文件（跳过 index.json 和 _scan_results.json）。"""
    samples: List[SampleFile] = []
    if not SAMPLES_DIR.exists():
        return samples
    for p in SAMPLES_DIR.iterdir():
        if not p.is_file():
            continue
        if p.name.startswith("_") or p.suffix == ".json":
            continue
        size = p.stat().st_size
        if size <= SIZE_TINY:
            bucket = "tiny"
        elif size <= SIZE_SMALL:
            bucket = "small"
        elif size <= SIZE_MEDIUM:
            bucket = "medium"
        else:
            bucket = "large"
        samples.append(SampleFile(path=p, name=p.name, size=size, size_bucket=bucket))
    return samples


def pick_samples(
    all_samples: List[SampleFile],
    count: int,
    bucket: str = "mixed",
    with_large_placeholder: bool = False,
) -> List[SampleFile]:
    """从样本池里挑 count 个文件。

    bucket: tiny / small / medium / large / mixed
    with_large_placeholder: 如果 bucket=large，返回一个 placeholder 让调用者生成大文件
    """
    if not all_samples and not with_large_placeholder:
        return []

    pool = all_samples
    if bucket != "mixed":
        pool = [s for s in all_samples if s.size_bucket == bucket]
        if not pool:
            # 如果指定 bucket 没文件，fallback 到 mixed
            pool = all_samples

    picked = random.choices(pool, k=count)
    return list(picked)


def make_large_temp_file(size_bytes: int = 5 * 1024 * 1024) -> Tuple[tempfile.NamedTemporaryFile, str]:
    """生成大文件（5MB），用于 large bucket 测试。"""
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".bin")
    tmp.write(os.urandom(size_bytes))
    tmp.flush()
    return tmp, tmp.name


# ---------------------------------------------------------------------------
# 核心压测客户端
# ---------------------------------------------------------------------------

class BenchmarkClient:
    def __init__(
        self,
        base_url: str,
        username: str = "admin",
        password: str = "admin123",
        timeout: float = DEFAULT_TIMEOUT,
    ) -> None:
        self.base_url = base_url.rstrip("/")
        self.username = username
        self.password = password
        self.timeout = timeout
        self.cookies: httpx.Cookies = httpx.Cookies()

    async def login(self) -> None:
        """登录 console，获取 session cookie。"""
        async with httpx.AsyncClient(timeout=self.timeout, verify=False, follow_redirects=True) as c:
            # 先 GET 登录页拿 csrf（如果有）
            await c.get(f"{self.base_url}/console/login")
            r = await c.post(
                f"{self.base_url}/console/login",
                data={"username": self.username, "password": self.password},
                headers={"Content-Type": "application/x-www-form-urlencoded"},
            )
            self.cookies = c.cookies
            if r.status_code not in (200, 302):
                raise RuntimeError(f"登录失败: HTTP {r.status_code}")

    async def upload_one(
        self,
        client: httpx.AsyncClient,
        files: List[Tuple[str, bytes, str]],  # [(filename, content_bytes, content_type), ...]
    ) -> RequestResult:
        """一次 upload 请求（可带多文件）。"""
        idx = random.randint(0, 999999)
        t0 = time.perf_counter()
        file_handles = []
        try:
            # 构造 files 参数
            upload_files = [("files", (name, content, ct)) for name, content, ct in files]
            r = await client.post(
                f"{self.base_url}/console/test/upload",
                files=upload_files,
            )
            t1 = time.perf_counter()
            ok = r.status_code == 200
            error = None
            files_count = len(files)
            if ok:
                try:
                    body = r.json()
                    if isinstance(body, dict):
                        results_list = body.get("results", [])
                        ok = all(it.get("verdict") not in ("error",) for it in results_list)
                except Exception as exc:
                    ok = False
                    error = f"JSON decode error: {exc}"
            else:
                error = f"HTTP {r.status_code}"
            avg_file_ms = int((t1 - t0) * 1000 / max(files_count, 1))
            return RequestResult(
                idx=idx,
                start_ts=t0,
                end_ts=t1,
                duration_ms=int((t1 - t0) * 1000),
                status_code=r.status_code,
                ok=ok,
                error=error,
                files_count=files_count,
                avg_file_ms=avg_file_ms,
            )
        except Exception as exc:  # noqa: BLE001
            t1 = time.perf_counter()
            return RequestResult(
                idx=idx,
                start_ts=t0,
                end_ts=t1,
                duration_ms=int((t1 - t0) * 1000),
                status_code=0,
                ok=False,
                error=str(exc),
                files_count=len(files),
            )
        finally:
            for fh in file_handles:
                try:
                    fh.close()
                except Exception:
                    pass

    async def run_scenario(
        self,
        name: str,
        concurrent: int,
        total_requests: int,
        files_per_request: int,
        size_bucket: str,
        samples: List[SampleFile],
    ) -> ScenarioResult:
        """跑一个场景。

        concurrent: 并发度（同时在飞的请求数上限）
        total_requests: 总共发多少请求
        files_per_request: 每个请求带几个文件（1=单文件场景，>1=批量）
        size_bucket: tiny/small/medium/large/mixed
        samples: 样本池
        """
        result = ScenarioResult(
            name=name,
            concurrent=concurrent,
            files_per_request=files_per_request,
            size_bucket=size_bucket,
            started_at=datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            finished_at="",
            total_duration_ms=0,
        )

        sem = asyncio.Semaphore(concurrent)

        async def _worker(idx: int, client: httpx.AsyncClient) -> RequestResult:
            async with sem:
                return await self._make_request(client, files_per_request, size_bucket, samples)

        connector = httpx.Limits(max_connections=concurrent * 2, max_keepalive_connections=concurrent)
        async with httpx.AsyncClient(
            timeout=self.timeout,
            verify=False,
            follow_redirects=True,
            cookies=self.cookies,
            limits=connector,
        ) as client:
            # 预热（不计入）
            if WARMUP_COUNT > 0:
                warmup_samples = pick_samples(samples, 1, size_bucket)
                for _ in range(min(WARMUP_COUNT, total_requests)):
                    try:
                        warmup_files = [(s.name, s.path.read_bytes(), "application/octet-stream") for s in warmup_samples]
                        await self.upload_one(client, warmup_files)
                    except Exception:
                        pass

            # 正式跑
            t_start = time.perf_counter()
            tasks = [asyncio.create_task(_worker(i, client)) for i in range(total_requests)]
            for fut in asyncio.as_completed(tasks):
                r = await fut
                result.results.append(r)
            t_end = time.perf_counter()

            result.total_duration_ms = int((t_end - t_start) * 1000)
            result.finished_at = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

        return result

    async def _make_request(
        self,
        client: httpx.AsyncClient,
        files_per_request: int,
        size_bucket: str,
        samples: List[SampleFile],
    ) -> RequestResult:
        """构造一个请求的文件列表，处理大文件临时生成。"""
        # 决定这个请求里每个文件的 bucket
        if size_bucket == "mixed":
            # 混合：随机挑 30% 小、50% 中、20% 大（大的现场生成）
            pass

        picked = pick_samples(samples, files_per_request, size_bucket)

        upload_files: List[Tuple[str, bytes, str]] = []
        temp_files_to_clean: List[str] = []

        for s in picked:
            if size_bucket == "large" or (size_bucket == "mixed" and random.random() < 0.2 and s.size_bucket != "large"):
                # 生成大文件 placeholder（用 .bin 模拟，如果没有 large bucket 样本）
                tmp_path = None
                large_candidates = [ps for ps in picked if ps.size_bucket == "large"]
                if large_candidates:
                    real = random.choice(large_candidates)
                    upload_files.append((real.name, real.path.read_bytes(), "application/octet-stream"))
                else:
                    tmp, tmp_path = make_large_temp_file(5 * 1024 * 1024)
                    upload_files.append((f"large_test_{tmp.name}.bin", Path(tmp.name).read_bytes(), "application/octet-stream"))
                    temp_files_to_clean.append(tmp.name)
                    tmp.close()
            else:
                upload_files.append((s.name, s.path.read_bytes(), "application/octet-stream"))

        try:
            return await self.upload_one(client, upload_files)
        finally:
            for tf in temp_files_to_clean:
                try:
                    os.unlink(tf)
                except Exception:
                    pass


# ---------------------------------------------------------------------------
# 报告格式化
# ---------------------------------------------------------------------------

def print_report(scenarios: List[ScenarioResult], base_url: str, workers: int = 4) -> None:
    """终端输出汇总报告。"""
    print("\n" + "=" * 80)
    print("  MeshSafe 并发压测报告")
    print("=" * 80)
    print(f"  目标:          {base_url}")
    print(f"  时间:          {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"  uvicorn:       --workers {workers}")
    print(f"  DB pool:       pool_size=20, max_overflow=30 (总上限 50)")
    print(f"  预热:          {WARMUP_COUNT} 次不计入")
    print("-" * 80)

    for sc in scenarios:
        print(f"\n【{sc.name}】")
        print(f"  并发数:          {sc.concurrent}")
        print(f"  请求文件数:      {sc.files_per_request} 个/请求")
        print(f"  文件大小档:      {sc.size_bucket}")
        print(f"  总请求数:        {sc.total_requests}")
        print(f"  总文件数:        {sc.total_files}")
        print(f"  总耗时:          {sc.total_duration_ms} ms  ({sc.total_duration_ms/1000:.2f}s)")
        print(f"  QPS:             {sc.qps:.2f} req/s")
        print(f"  文件吞吐:        {sc.files_per_sec:.2f} files/s")
        print(f"  单请求耗时:")
        print(f"    平均:          {sc.avg_ms:.0f} ms")
        print(f"    P50:           {sc.p50_ms:.0f} ms")
        print(f"    P95:           {sc.p95_ms:.0f} ms")
        print(f"    P99:           {sc.p99_ms:.0f} ms")
        print(f"  错误率:          {sc.error_rate*100:.1f}%")
        if sc.slowest:
            slowest_msg = sc.slowest.error or "ok"
            print(f"  最慢请求:        {sc.slowest.duration_ms} ms  ({slowest_msg[:60]})")

    # 瓶颈分析
    print("\n" + "-" * 80)
    print("【瓶颈分析】")
    _analyze(scenarios)

    print("\n" + "=" * 80)


def _analyze(scenarios: List[ScenarioResult]) -> None:
    if not scenarios:
        return
    for sc in scenarios:
        if sc.concurrent == 1 and sc.files_per_request == 1 and sc.size_bucket in ("mixed", "medium"):
            baseline = sc.avg_ms
            print(f"  基准单文件串行:  {baseline:.0f} ms")

    # 并发场景看扩展效率
    for sc in scenarios:
        if sc.concurrent > 1 and sc.files_per_request == 1:
            print(f"  {sc.concurrent} 并发单文件:")
            if sc.avg_ms < 300:
                print(f"    ✅ 响应快（avg={sc.avg_ms:.0f}ms）")
            elif sc.avg_ms < 800:
                print(f"    ⚠️ 中等（avg={sc.avg_ms:.0f}ms）")
            else:
                print(f"    🔴 较慢（avg={sc.avg_ms:.0f}ms）—— 可能 uvicorn worker 不够或 DB pool 满")

            if sc.error_rate > 0:
                print(f"    🔴 错误率 {sc.error_rate*100:.1f}% —— 检查 DB 连接池 / ClamAV / YARA")

        if sc.files_per_request > 1:
            per_file = sc.avg_ms / sc.files_per_request
            print(f"  批量 {sc.files_per_request} 文件·单用户:")
            print(f"    总耗时 ≈ {sc.avg_ms:.0f}ms  单文件均摊 ≈ {per_file:.0f}ms")
            if per_file > 300:
                print(f"    ⚠️ 批量内部 for 循环串行，单文件均摊偏高—— 可考虑 ThreadPoolExecutor 并行化")
            else:
                print(f"    ✅ 批量表现正常")


def to_json_report(scenarios: List[ScenarioResult], base_url: str) -> dict:
    return {
        "generated_at": datetime.now().isoformat(),
        "base_url": base_url,
        "uvicorn_workers": 4,
        "db_pool": {"pool_size": 20, "max_overflow": 30},
        "scenarios": [
            {
                "name": sc.name,
                "concurrent": sc.concurrent,
                "files_per_request": sc.files_per_request,
                "size_bucket": sc.size_bucket,
                "started_at": sc.started_at,
                "finished_at": sc.finished_at,
                "total_duration_ms": sc.total_duration_ms,
                "total_requests": sc.total_requests,
                "total_files": sc.total_files,
                "qps": round(sc.qps, 2),
                "files_per_sec": round(sc.files_per_sec, 2),
                "latency_ms": {
                    "avg": round(sc.avg_ms, 1),
                    "p50": round(sc.p50_ms, 1),
                    "p95": round(sc.p95_ms, 1),
                    "p99": round(sc.p99_ms, 1),
                },
                "error_rate": round(sc.error_rate, 4),
                "slowest_ms": sc.slowest.duration_ms if sc.slowest else None,
            }
            for sc in scenarios
        ],
    }


# ---------------------------------------------------------------------------
# 主入口
# ---------------------------------------------------------------------------

async def _main() -> None:
    parser = argparse.ArgumentParser(description="MeshSafe 并发压测")
    parser.add_argument("--base-url", default=DEFAULT_BASE_URL, help="MeshSafe 基地址")
    parser.add_argument("--username", default="admin")
    parser.add_argument("--password", default="admin123")
    parser.add_argument(
        "--concurrent",
        nargs="+",
        type=int,
        default=[50, 100],
        help="并发度档位（可多值）",
    )
    parser.add_argument(
        "--batch-sizes",
        nargs="+",
        type=int,
        default=[10, 50],
        help="批量上传文件数（单用户·多文件场景）",
    )
    parser.add_argument(
        "--requests-per-batch",
        type=int,
        default=5,
        help="批量场景下发几个请求（每个请求带 N 文件）",
    )
    parser.add_argument(
        "--requests-per-concurrent",
        type=int,
        default=2,
        help="并发场景下每个并发度发几个请求",
    )
    parser.add_argument("--timeout", type=float, default=DEFAULT_TIMEOUT)
    parser.add_argument("--report", default=None, help="JSON 报告输出路径")
    args = parser.parse_args()

    # 1. 加载样本
    all_samples = load_samples()
    if not all_samples:
        print(f"❌ tests/samples/ 下没找到样本文件")
        sys.exit(1)
    sizes = sorted([(s.size, s.size_bucket) for s in all_samples], reverse=True)
    print(f"📦 样本池: {len(all_samples)} 个文件")
    for bucket in ("tiny", "small", "medium", "large"):
        count = sum(1 for _, b in sizes if b == bucket)
        if count:
            print(f"   {bucket:7s}: {count} 个")

    # 2. 登录
    print(f"\n🔐 登录 {args.base_url} ...")
    client = BenchmarkClient(base_url=args.base_url, username=args.username, password=args.password, timeout=args.timeout)
    await client.login()
    print("   ✅ 登录成功")

    # 3. 跑场景
    scenarios: List[ScenarioResult] = []

    # 场景 A：基准串行（1 并发）
    print("\n" + "=" * 80)
    print("【场景 A】基准串行·单文件·混合大小")
    print("=" * 80)
    sc_a = await client.run_scenario(
        name="A1. 基准串行·单文件",
        concurrent=1,
        total_requests=20,
        files_per_request=1,
        size_bucket="mixed",
        samples=all_samples,
    )
    scenarios.append(sc_a)

    # 场景 B：多用户并发·单文件
    for conc in args.concurrent:
        print(f"\n{'='*80}")
        print(f"【场景 B】{conc} 并发·单文件·混合大小")
        print("=" * 80)
        sc = await client.run_scenario(
            name=f"B.{conc}·并发单文件·混合",
            concurrent=conc,
            total_requests=conc * args.requests_per_concurrent,
            files_per_request=1,
            size_bucket="mixed",
            samples=all_samples,
        )
        scenarios.append(sc)

    # 场景 C：单用户·批量多文件
    for batch_n in args.batch_sizes:
        print(f"\n{'='*80}")
        print(f"【场景 C】单用户·{batch_n} 文件批量上传")
        print("=" * 80)
        sc = await client.run_scenario(
            name=f"C.{batch_n}·单用户批量",
            concurrent=1,
            total_requests=args.requests_per_batch,
            files_per_request=batch_n,
            size_bucket="mixed",
            samples=all_samples,
        )
        scenarios.append(sc)

    # 场景 D：多用户并发·批量
    conc_mix = max(10, min(args.concurrent[0] // 5, 20))
    batch_mix = 10
    print(f"\n{'='*80}")
    print(f"【场景 D】{conc_mix} 并发·{batch_mix} 文件/请求（最坏情况）")
    print("=" * 80)
    sc_d = await client.run_scenario(
        name=f"D.{conc_mix}并发×{batch_mix}文件",
        concurrent=conc_mix,
        total_requests=conc_mix,
        files_per_request=batch_mix,
        size_bucket="mixed",
        samples=all_samples,
    )
    scenarios.append(sc_d)

    # 4. 输出报告
    print_report(scenarios, args.base_url)

    # 5. JSON 报告
    if args.report:
        report = to_json_report(scenarios, args.base_url)
        Path(args.report).parent.mkdir(parents=True, exist_ok=True)
        Path(args.report).write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"\n📄 JSON 报告已保存: {args.report}")


if __name__ == "__main__":
    asyncio.run(_main())