# -*- coding: utf-8 -*-
"""LLM 路由 / DLP NFR 压测（Issue #88，对应 PRD 第 9 章）。

验证目标：
- DLP 敏感数据拦截率 **100%**（目标 100%，PRD 5.4 / 9 章）；
- 敏感度路由敏感命中率 ≥ **96.5%**（PRD 9 章路由基线）；
- 压测吞吐与时延在可接受区间。

用合成数据集对 DlpEngine / SensitivityRouter 做规模化压测，输出 SLA 达标结论。
"""
from __future__ import annotations

# 引导加载内核模块
import tests.perf._bootstrap  # noqa: F401

import random
import time
import unittest

from llm_gateway import DLP_DEFAULT_RULES, DlpEngine, RouteTarget, SensitivityRouter


# ---------------------------------------------------------------------------
# 合成数据集（可重复：固定随机种子）
# ---------------------------------------------------------------------------
_RANDOM = random.Random(20260804)

# 敏感样本池（应被 DLP 100% 拦截 / 路由到 local 或 block）
SENSITIVE_SAMPLES = [
    # PII：身份证号
    "员工 {n} 的身份证 110101199003078834 工资明细",
    # PII：手机号
    "联系工程师手机 13912345678 确认 {n} 号炉工况",
    # PII：邮箱
    "报告发送至 ops_{n}@xayunmei.com 并抄送管理",
    # 凭证：访问密钥
    "使用 AKIAIOSFODNN7EXAMPLE2010 上传 {n} 号数据",
    # 凭证：键值对口令
    "数据库 password=Admin@{n} 请勿外泄",
    # 工艺敏感词
    "{n} 号氯化炉核心配方 已更新，请审核",
    # IP（第四段用 {ip} 占位，映射为 0~255 合法段，避免序号破坏格式）
    "PLC 控制器地址 192.168.10.{ip} 离线告警（节点 {n}）",
]

# 安全（应被路由 block）
BLOCK_SAMPLES = [
    "立即执行 {n} 号炉停机 操作",
    "{n} 号机组紧急停机 检修",
]

# 清洁样本池（应被 DLP 放行 / 路由 local 默认）
CLEAN_SAMPLES = [
    "{n} 号氯化炉当前炉温 880℃ 是否在正常区间",
    "查询 {n} 号炉氯气流量历史趋势",
    "{n} 号冷却水进出口温度差报警阈值是多少",
    "展示 {n} 号炉今日产量统计报表",
    "{n} 号塔顶压力与塔底压力压差分析",
]


def _expand(pool, n):
    """从模板池扩样为 n 条（带编号，保证多样性）。

    注：模板中以 `{ip}` 占位的字段会映射为合法 IPv4 第四段（0~255 取模），
    避免序号破坏 PII 格式导致规则漏匹配（压测验证的是规则本身而非数据生成）。
    """
    out = []
    for i in range(n):
        tpl = pool[i % len(pool)]
        out.append(tpl.format(n=(i + 1), ip=(i % 256)))
    return out


class DlpInterceptNFRTest(unittest.TestCase):
    """DLP 敏感数据拦截压测（目标 100% 拦截）。"""

    @classmethod
    def setUpClass(cls) -> None:
        cls.engine = DlpEngine(rules=list(DLP_DEFAULT_RULES), fail_closed=True)
        # 规模：敏感样本 2000 条（PII / 凭证 / 工艺敏感词全覆盖）
        cls.sensitive = _expand(SENSITIVE_SAMPLES, 2000)

    def test_dlp_intercept_rate_100_percent(self) -> None:
        """压测：2000 条敏感样本，DLP 拦截率必须 = 100%。"""
        blocked = 0
        t0 = time.perf_counter()
        for sample in self.sensitive:
            if self.engine.check_outbound({"query": sample}).blocked:
                blocked += 1
        elapsed = time.perf_counter() - t0

        total = len(self.sensitive)
        rate = blocked / total
        # PRD 目标 100% 拦截：必须严格达标
        self.assertEqual(rate, 1.0,
                         f"DLP 拦截率 {rate:.4%} 未达 100% 目标（{total-blocked} 条漏拦）")
        # 吞吐与时延（信息性断言：单条 < 5ms，整体 ≥ 500 qps）
        avg_ms = elapsed / total * 1000
        qps = total / elapsed
        self.assertLess(avg_ms, 5.0, f"DLP 单条均时延 {avg_ms:.3f}ms 超 5ms")
        print(f"\n[DLP 压测] 样本 {total} 条，拦截 {blocked}，"
              f"拦截率 100%，均时延 {avg_ms:.3f}ms，吞吐 {qps:.0f} qps")

    def test_dlp_clean_passthrough(self) -> None:
        """清洁样本不被误拦（放行率应 = 100%，零误报）。"""
        clean = _expand(CLEAN_SAMPLES, 1000)
        allowed = sum(
            1 for s in clean
            if not self.engine.check_outbound({"query": s}).blocked
        )
        self.assertEqual(allowed, len(clean),
                         "清洁样本被误拦，DLP 误报率 > 0")


class RouterSensitivityNFRTest(unittest.TestCase):
    """敏感度路由命中率压测（目标 ≥ 96.5%）。"""

    @classmethod
    def setUpClass(cls) -> None:
        cls.router = SensitivityRouter()
        cls.sensitive = _expand(SENSITIVE_SAMPLES, 1500)
        cls.block = _expand(BLOCK_SAMPLES, 500)

    def test_router_sensitive_never_cloud(self) -> None:
        """敏感样本（PII/凭证/工艺词）绝不路由到 cloud（命中率 100%）。"""
        leaked = 0
        for sample in self.sensitive:
            # 先 DLP 检查，模拟网关真实流程：DLP 拦截即 block
            dlp = DlpEngine(rules=list(DLP_DEFAULT_RULES))
            blocked = dlp.check_outbound({"query": sample}).blocked
            decision = self.router.route(sample, dlp_blocked=blocked)
            if decision.target == RouteTarget.CLOUD:
                leaked += 1
        rate = 1 - leaked / len(self.sensitive)
        # PRD 路由基线 ≥ 96.5%；这里敏感样本应 100% 不入云端
        self.assertGreaterEqual(rate, 0.965,
                                f"敏感样本路由泄漏率 {(1-rate):.4%}，未达 ≥96.5% 命中基线")
        print(f"\n[路由压测] 敏感样本 {len(self.sensitive)} 条，"
              f"未泄漏到云端 {len(self.sensitive)-leaked}，命中率 {rate:.2%}")

    def test_router_block_safety_commands(self) -> None:
        """安全指令（停机）必须路由 block（转人工）。"""
        missed = 0
        for sample in self.block:
            decision = self.router.route(sample)
            if decision.target != RouteTarget.BLOCK:
                missed += 1
        self.assertEqual(missed, 0,
                         f"{missed} 条安全指令未路由 block（应 100% 转人工）")


if __name__ == "__main__":
    unittest.main()