# -*- coding: utf-8 -*-
"""DLP 敏感数据拦截引擎（dlp）单元测试：规则匹配 / 出站拦截 / 审计 / 模板加载。"""
import os
import sys
import unittest

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import _bootstrap  # noqa: F401

from llm_gateway.dlp import (  # noqa: E402
    DLP_DEFAULT_RULES,
    DlpEngine,
    DlpRule,
    DlpRuleKind,
)

CONFIG_PATH = os.path.join(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
    "config", "dlp.template.yaml",
)


class BuiltinRulesTest(unittest.TestCase):
    """内置保底规则（无任何配置时通用 PII 也默认拦截）。"""

    def setUp(self):
        self.engine = DlpEngine()

    def test_default_rules_present(self):
        self.assertGreater(len(DLP_DEFAULT_RULES), 0)
        self.assertGreaterEqual(self.engine.rule_count, len(DLP_DEFAULT_RULES))

    def test_id_card_blocked(self):
        result = self.engine.check_outbound({"query": "员工身份证 110101199003071234 入职"})
        self.assertTrue(result.blocked)
        self.assertEqual(result.reason, "hit")
        self.assertTrue(any(h.category == "pii" for h in result.hits))

    def test_mobile_blocked(self):
        result = self.engine.check_outbound({"query": "联系 13812345678"})
        self.assertTrue(result.blocked)
        self.assertIn("pii_mobile", {h.rule for h in result.hits})

    def test_cloud_ak_blocked(self):
        result = self.engine.check_outbound({"output": "ak = AKIAIOSFODNN7EXAMPLE"})
        self.assertTrue(result.blocked)
        self.assertIn("credential_ak", {h.rule for h in result.hits})

    def test_clean_content_allowed(self):
        result = self.engine.check_outbound({
            "query": "今天炉况怎么样",
            "context": "炉温控制在 850-950°C 范围内。",
            "output": "当前运行平稳。",
        })
        self.assertFalse(result.blocked)
        self.assertEqual(result.decision, "allow")


class KeywordRuleTest(unittest.TestCase):
    def test_cn_keyword_hit(self):
        engine = DlpEngine(rules=[
            DlpRule(name="proc_formula", category="process-keyword",
                    kind=DlpRuleKind.KEYWORD, pattern="配方"),
        ])
        result = engine.check_outbound({"output": "已按配方调整投料"})
        self.assertTrue(result.blocked)
        self.assertIn("proc_formula", {h.rule for h in result.hits})
        hit = next(h for h in result.hits if h.rule == "proc_formula")
        self.assertEqual(hit.matched, "配方")

    def test_en_keyword_case_insensitive(self):
        engine = DlpEngine(rules=[
            DlpRule(name="secret_word", category="credential",
                    kind=DlpRuleKind.KEYWORD, pattern="SecretKey"),
        ])
        self.assertTrue(engine.check_outbound({"output": "the secretkey is hidden"}).blocked)
        self.assertTrue(engine.check_outbound({"output": "the SECRETKEY is hidden"}).blocked)

    def test_no_match_allowed(self):
        engine = DlpEngine(rules=[
            DlpRule(name="proc_formula", category="process-keyword",
                    kind=DlpRuleKind.KEYWORD, pattern="配方"),
        ])
        result = engine.check_outbound({"output": "炉温正常"})
        self.assertFalse(result.blocked)


class RegexRuleTest(unittest.TestCase):
    def test_structured_formula_ratio(self):
        engine = DlpEngine(rules=[
            DlpRule(name="proc_ratio", category="process-parameter",
                    kind=DlpRuleKind.REGEX, pattern=r"配比\s*[:：]?\s*[\d.]+%?"),
        ])
        result = engine.check_outbound({"output": "当前配比：0.35 保持不变"})
        self.assertTrue(result.blocked)
        self.assertEqual(result.hits[0].matched, "配比：0.35")


class OutboundCheckTest(unittest.TestCase):
    """出站检查：多 part 聚合 / 命中定位 / 拦截原因。"""

    def setUp(self):
        self.engine = DlpEngine()

    def test_hit_in_context_blocks_whole_packet(self):
        result = self.engine.check_outbound({
            "query": "请问这个员工怎么样",
            "context": "档案：13900000000",
            "output": "该员工表现良好。",
        })
        self.assertTrue(result.blocked)
        # 命中定位在 context part
        self.assertEqual(result.hits[0].part, "context")

    def test_empty_parts_allowed(self):
        self.assertFalse(self.engine.check_outbound({}).blocked)
        self.assertFalse(self.engine.check_outbound({"query": ""}).blocked)


class FailClosedTest(unittest.TestCase):
    def test_no_rules_fail_closed_blocks(self):
        engine = DlpEngine(rules=[], include_builtins=False, fail_closed=True)
        result = engine.check_outbound({"query": "任意内容"})
        self.assertTrue(result.blocked)
        self.assertEqual(result.reason, "no_rules")

    def test_fail_open_without_rules_allows(self):
        engine = DlpEngine(rules=[], include_builtins=False, fail_closed=False)
        self.assertFalse(engine.check_outbound({"query": "任意内容"}).blocked)


class MaskTest(unittest.TestCase):
    def test_mask_redacts_hits(self):
        engine = DlpEngine()
        masked = engine.mask("联系 13812345678 或 13900000000")
        self.assertNotIn("13812345678", masked)
        self.assertNotIn("13900000000", masked)
        self.assertIn("<pii>", masked)

    def test_mask_clean_text_unchanged(self):
        engine = DlpEngine()
        text = "炉温正常"
        self.assertEqual(engine.mask(text), text)


class AuditTest(unittest.TestCase):
    def test_audit_record_fields(self):
        engine = DlpEngine()
        result = engine.check_outbound({"query": "身份证 110101199003071234"})
        audit = result.audit
        self.assertEqual(audit["channel"], "outbound")
        self.assertEqual(audit["decision"], "block")
        self.assertEqual(audit["reason"], "hit")
        self.assertGreaterEqual(audit["hit_count"], 1)
        # 审计不落明文敏感内容，只记脱敏占位
        serialized = str(audit)
        self.assertNotIn("110101199003071234", serialized)
        self.assertIn("<pii>", serialized)

    def test_audit_buffer_and_hook(self):
        seen = []
        engine = DlpEngine(audit_hook=seen.append)
        engine.check_outbound({"query": "13812345678"})
        engine.check_outbound({"query": "炉温正常"})
        self.assertEqual(len(seen), 2)
        self.assertEqual(len(engine.drain_audit()), 2)
        self.assertEqual(engine.drain_audit(), [])


class TemplateConfigTest(unittest.TestCase):
    """从真实模板资产（config/dlp.template.yaml）加载。"""

    def test_load_example_assets(self):
        engine = DlpEngine.from_template_config(CONFIG_PATH)
        self.assertEqual(engine.template, "ti-cl4")
        self.assertEqual(engine.version, "1.0.0")
        self.assertTrue(engine.fail_closed)
        # 模板规则 + 内置保底规则
        names = set(engine.rule_names())
        self.assertIn("proc_cl2_flow", names)
        self.assertIn("proc_furnace_temp", names)
        self.assertIn("pii_mobile", names)  # 内置保底仍在

    def test_template_industry_keyword_blocks(self):
        engine = DlpEngine.from_template_config(CONFIG_PATH)
        result = engine.check_outbound({"output": "今日氯气流量为 5.2 t/h"})
        self.assertTrue(result.blocked)
        self.assertIn("proc_cl2_flow", {h.rule for h in result.hits})

    def test_template_regex_blocks(self):
        engine = DlpEngine.from_template_config(CONFIG_PATH)
        result = engine.check_outbound({"output": "配比：0.42 已确认"})
        self.assertTrue(result.blocked)
        self.assertIn("proc_formula_detail", {h.rule for h in result.hits})

    def test_template_rule_overrides_builtin(self):
        # 同名规则（pii_mobile）在模板中自定义 pattern 后覆盖内置
        engine = DlpEngine.from_template_config(CONFIG_PATH)
        result = engine.check_outbound({"query": "联系 13812345678"})
        # 内置手机号规则仍在（模板未覆盖），应照常命中
        self.assertTrue(result.blocked)


class TemplateOverrideTest(unittest.TestCase):
    def test_same_name_overrides_builtin(self):
        engine = DlpEngine(rules=[
            DlpRule(name="pii_mobile", category="pii",
                    kind=DlpRuleKind.REGEX, pattern=r"\b1[3-9]\d{9}\b"),
        ])
        names = [r.name for r in engine._rules]
        self.assertEqual(names.count("pii_mobile"), 1)


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