# -*- coding: utf-8 -*-
"""iAOP-Core · LLM 网关 —— DLP 敏感数据拦截（Issue #48，EPIC #6 子任务）。

对应 PRD 5.4「④ LLM 网关 + RAG」：
- 数据不出厂：敏感数据**本地闭环**，仅脱敏/公开内容可走云端 API；
- 验收：**DLP 拦截率 100%**，敏感词/DLP 规则作为模板配置点；
- 用户操作流程：用户提问 → 路由判断敏感级 → 本地/云端生成 →
  RAG 溯源校验 → 返回答案（出站前经 DLP 检查）。

本模块实现 **出站 DLP 防线**：任何要发往**云端**的内容
（用户 query / RAG 检索到的 context / 模型输出）先经 `DlpEngine.check_outbound`
检查，命中任一 block 规则即**整包拦截**（fail-closed，目标 100% 拦截），
并落**结构化审计日志**（全量，PRD NFR 安全-数据）。

规则来源：
- 模板资产 `config/dlp.template.yaml`（行业敏感词 + 通用 PII 正则），
  换行业只改该资产，内核零改动（对齐 rag-kb 模板化思想）；
- `DLP_DEFAULT_RULES` 内置保底规则：即使未加载任何配置，
  通用 PII（身份证/手机号/邮箱/IP/访问密钥等）也默认生效。

测试：`python -m unittest discover -s tests -v`（在 core/llm-gateway 目录下执行）。
"""
from __future__ import annotations

import re
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Callable, Dict, List, Optional, Tuple

# ---------------------------------------------------------------------------
# 轻量 YAML 子集解析（零第三方依赖，递归下降）：map / list / scalar / 注释。
# 与 rag-kb/templating.py 同款（模块内自持一份，保持模块零耦合）。
# 足以解析 `config/dlp.template.yaml` 这类模板资产。
# ---------------------------------------------------------------------------


def _parse_scalar(text: str) -> str:
    """去掉标量两侧引号与行内注释（`key: value # comment`）。"""
    t = text.split(" #", 1)[0].strip()
    if len(t) >= 2 and t[0] == t[-1] and t[0] in ("'", '"'):
        return t[1:-1]
    return t


def _strip_comments(lines: List[str]) -> List[Tuple[str, int]]:
    """剔除空行与整行注释，保留行号（1 起）用于报错定位。"""
    out = []
    for i, ln in enumerate(lines):
        s = ln.strip()
        if not s or s.startswith("#"):
            continue
        out.append((ln, i + 1))
    return out


def _parse_node(lines: List[Tuple[str, int]], i: int, indent: int):
    """递归解析从 lines[i] 开始、缩进为 `indent` 的一个节点。

    返回 `(value, next_i)`：value 为 dict / list / str，next_i 为下一个
    未消费行的下标。
    """
    text, no = lines[i]
    lead = len(text) - len(text.lstrip(" "))

    # ---- list 节点：`- item` 或 `- key: val`（map 项） ----
    if text.lstrip(" ").startswith("- "):
        items: List[object] = []
        while i < len(lines):
            t, no2 = lines[i]
            stripped = t.lstrip(" ")
            if not stripped.startswith("- "):
                break
            lead_j = len(t) - len(t.lstrip(" "))
            if lead_j != indent:
                break
            item_text = stripped[2:].strip()
            if not item_text:
                raise ValueError(f"dlp.yaml 第 {no2} 行：list 项为空")
            if ":" in item_text:
                map_indent = len(t) - len(t.lstrip(" ")) + 2
                lines[i] = (" " * map_indent + item_text, no2)
                v, i = _parse_node(lines, i, map_indent)
                items.append(v)
            else:
                items.append(_parse_scalar(item_text))
                i += 1
        return items, i

    # ---- map 节点：`key: value` / `key:`（嵌套值） ----
    result: Dict[str, object] = {}
    while i < len(lines):
        t, no = lines[i]
        lead_j = len(t) - len(t.lstrip(" "))
        if lead_j < indent or t.lstrip(" ").startswith("- "):
            break
        if lead_j > indent:
            raise ValueError(f"dlp.yaml 第 {no} 行缩进异常（期望 {indent}，实际 {lead_j}）")
        if ":" not in t:
            raise ValueError(f"dlp.yaml 第 {no} 行不是合法键值对：{t!r}")
        key, _, rest = t.partition(":")
        key = key.strip()
        rest = rest.strip()
        if rest:
            result[key] = _parse_scalar(rest)
            i += 1
            continue
        if i + 1 >= len(lines):
            raise ValueError(f"dlp.yaml 第 {no} 行 {key!r} 缺少值")
        sub_indent = len(lines[i + 1][0]) - len(lines[i + 1][0].lstrip(" "))
        if sub_indent <= indent:
            raise ValueError(f"dlp.yaml 第 {no} 行 {key!r} 缺少值（无嵌套内容）")
        v, i = _parse_node(lines, i + 1, sub_indent)
        result[key] = v
    return result, i


def _load_yaml_text(text: str) -> Dict[str, object]:
    """解析 YAML 子集 → 嵌套 dict/list。顶层必须为 map。"""
    lines = _strip_comments(text.splitlines())
    if not lines:
        return {}
    top_indent = len(lines[0][0]) - len(lines[0][0].lstrip(" "))
    value, next_i = _parse_node(lines, 0, top_indent)
    if not isinstance(value, dict):
        raise ValueError("dlp.yaml 顶层必须是 map")
    if next_i < len(lines):
        raise ValueError(
            f"dlp.yaml 第 {lines[next_i][1]} 行：顶层存在多个节点（缩进不一致）"
        )
    return value


# ---------------------------------------------------------------------------
# 规则模型
# ---------------------------------------------------------------------------


class DlpRuleKind:
    """规则匹配方式：`keyword` 子串匹配 / `regex` 正则匹配。"""

    KEYWORD = "keyword"
    REGEX = "regex"


@dataclass(frozen=True)
class DlpRule:
    """一条 DLP 敏感数据规则（来自模板资产或内置保底规则）。

    - `name`        规则唯一名（审计用）；
    - `category`    敏感类别（pii / process-keyword / credential 等）；
    - `kind`        keyword 子串匹配 或 regex 正则匹配；
    - `pattern`     关键词原文，或正则表达式；
    - `description` 规则说明（人读）。
    """

    name: str
    category: str
    pattern: str
    kind: str = DlpRuleKind.KEYWORD
    description: str = ""

    def _compiled(self) -> "re.Pattern[str]":
        if self.kind == DlpRuleKind.REGEX:
            return re.compile(self.pattern, re.IGNORECASE)
        return re.compile(re.escape(self.pattern), re.IGNORECASE)

    def find(self, text: str) -> List[Tuple[str, int, int]]:
        """返回文本中所有命中片段 `(matched, start, end)`（不重叠）。"""
        out: List[Tuple[str, int, int]] = []
        if not text:
            return out
        if self.kind == DlpRuleKind.REGEX:
            for m in self._compiled().finditer(text):
                out.append((m.group(0), m.start(), m.end()))
        else:
            # 大小写不敏感子串匹配（中文无大小写，直接命中）
            lower = text.lower()
            needle = self.pattern.lower()
            start = 0
            while True:
                pos = lower.find(needle, start)
                if pos < 0:
                    break
                out.append((text[pos:pos + len(self.pattern)], pos, pos + len(self.pattern)))
                start = pos + max(len(self.pattern), 1)
        return out


@dataclass(frozen=True)
class DlpHit:
    """一次敏感命中：哪条规则、哪段文本、位于何处（可选所属出站 part）。"""

    rule: str
    category: str
    matched: str
    start: int
    end: int
    part: str = ""

    def to_dict(self) -> Dict[str, object]:
        # 审计只记录脱敏预览（`<category>` 占位）与所属 part，不落明文敏感内容
        return {
            "rule": self.rule,
            "category": self.category,
            "part": self.part,
            "matched": f"<{self.category}>",
            "span": [self.start, self.end],
        }


@dataclass(frozen=True)
class DlpResult:
    """一次出站检查的结论。

    - `decision`   allow（放行）| block（拦截）；
    - `reason`     拦截原因（`hit` 命中敏感规则 / `no_rules` 未配置规则且 fail-closed）；
    - `hits`       命中明细；
    - `audit`      本次检查的结构化审计记录（JSON 序列化友好）。
    """

    decision: str
    reason: str
    hits: List[DlpHit]
    audit: Dict[str, object]

    @property
    def blocked(self) -> bool:
        return self.decision == "block"


# ---------------------------------------------------------------------------
# 内置保底规则：未加载任何模板配置时，通用 PII 也默认生效（目标 100% 拦截）
# ---------------------------------------------------------------------------

DLP_DEFAULT_RULES: List[DlpRule] = [
    # ---- PII / 凭证（通用，任何行业模板都适用） ----
    DlpRule(name="pii_id_card", category="pii",
            kind=DlpRuleKind.REGEX, pattern=r"\b\d{17}[\dXx]\b",
            description="身份证号（18 位）"),
    DlpRule(name="pii_mobile", category="pii",
            kind=DlpRuleKind.REGEX, pattern=r"\b1[3-9]\d{9}\b",
            description="中国大陆手机号"),
    DlpRule(name="pii_email", category="pii",
            kind=DlpRuleKind.REGEX, pattern=r"[\w.+-]+@[\w-]+\.[\w.-]+",
            description="邮箱地址"),
    DlpRule(name="pii_ipv4", category="pii",
            kind=DlpRuleKind.REGEX, pattern=r"\b(?:\d{1,3}\.){3}\d{1,3}\b",
            description="IPv4 地址"),
    DlpRule(name="credential_ak", category="credential",
            kind=DlpRuleKind.REGEX, pattern=r"\b(?:AKIA|LTAI)[0-9A-Z]{16,}\b",
            description="云访问密钥 ID（AWS AKIA / 阿里云 LTAI 样式）"),
    DlpRule(name="credential_kv", category="credential",
            kind=DlpRuleKind.REGEX,
            pattern=r"(?:access[_-]?key|secret|password|token)\s*[:=]\s*\S+",
            description="密钥/口令键值对"),
    # ---- 工艺敏感词示例（行业模板会按需增删，见 config/dlp.template.yaml） ----
    DlpRule(name="process_formula", category="process-keyword",
            kind=DlpRuleKind.KEYWORD, pattern="配方",
            description="配方类敏感词"),
]


class DlpEngine:
    """出站 DLP 检查引擎（fail-closed：命中即拦截，未配置规则即拦截）。

    用法：
    ```python
    engine = DlpEngine.from_template_config("config/dlp.template.yaml")
    result = engine.check_outbound({
        "query":   user_query,
        "context": rag_context,
        "output":  llm_answer,
    })
    if result.blocked:
        # 不发往云端：转本地处理 / 转人工 / 拒绝（PRD 5.4）
        ...
    ```
    """

    def __init__(self,
                 rules: Optional[List[DlpRule]] = None,
                 fail_closed: bool = True,
                 template: str = "builtin",
                 version: str = "1.0.0",
                 include_builtins: bool = True,
                 audit_hook: Optional[Callable[[Dict[str, object]], None]] = None):
        # 内置保底规则 + 模板规则（模板同名规则覆盖内置，实现模板定制）
        merged: Dict[str, DlpRule] = {}
        if include_builtins:
            merged.update({r.name: r for r in DLP_DEFAULT_RULES})
        for r in rules or []:
            merged[r.name] = r
        self._rules: List[DlpRule] = list(merged.values())
        self.fail_closed = fail_closed
        self.template = template
        self.version = version
        self.audit_hook = audit_hook
        self._audit_buffer: List[Dict[str, object]] = []

    # ------------------------------------------------------------------
    # 构造
    # ------------------------------------------------------------------

    @classmethod
    def from_template_config(cls, path: str,
                             audit_hook: Optional[Callable[[Dict[str, object]], None]] = None
                             ) -> "DlpEngine":
        """从模板 DLP 规则资产加载引擎（期望结构见 `config/dlp.template.yaml`）。"""
        with open(path, "r", encoding="utf-8") as fh:
            data = _load_yaml_text(fh.read())

        template = str(data.get("template", "")).strip()
        version = str(data.get("version", "1.0.0")).strip() or "1.0.0"
        mode = str(data.get("mode", "fail-closed")).strip().lower()
        fail_closed = mode != "fail-open"

        rules: List[DlpRule] = []
        raw_rules = data.get("rules") or []
        if not isinstance(raw_rules, list):
            raise ValueError("dlp.yaml rules 必须是 list")
        for item in raw_rules:
            if not isinstance(item, dict):
                raise ValueError(f"dlp.yaml rules 项必须是 map，实际 {item!r}")
            name = str(item.get("name", "")).strip()
            if not name:
                raise ValueError("dlp.yaml 规则缺少 name")
            kind = str(item.get("kind", "keyword")).strip().lower()
            if kind not in (DlpRuleKind.KEYWORD, DlpRuleKind.REGEX):
                raise ValueError(f"dlp.yaml 规则 {name!r} 的 kind 非法：{kind!r}")
            pattern = str(item.get("pattern", "")).strip()
            if not pattern:
                raise ValueError(f"dlp.yaml 规则 {name!r} 缺少 pattern")
            rules.append(DlpRule(
                name=name,
                category=str(item.get("category", "custom")).strip() or "custom",
                kind=kind,
                pattern=pattern,
                description=str(item.get("description", "")).strip(),
            ))
        return cls(rules=rules, fail_closed=fail_closed,
                   template=template, version=version, audit_hook=audit_hook)
    # ------------------------------------------------------------------
    # 检查
    # ------------------------------------------------------------------

    def inspect(self, text: str) -> List[DlpHit]:
        """对单段文本做敏感规则扫描，返回全部命中（按文本位置排序）。"""
        hits: List[DlpHit] = []
        for rule in self._rules:
            for matched, start, end in rule.find(text or ""):
                hits.append(DlpHit(
                    rule=rule.name, category=rule.category,
                    matched=matched, start=start, end=end,
                ))
        hits.sort(key=lambda h: (h.start, h.end))
        return hits

    def mask(self, text: str) -> str:
        """把命中片段替换为 `<category>` 占位（供"脱敏后可出站"的后续流程）。

        注意：本任务（Issue #48）的出站策略为**命中即拦截**，不自动放行
        脱敏内容；mask 仅作为审计预览与后续脱敏通道的辅助能力。
        """
        hits = self.inspect(text)
        if not hits:
            return text
        out = []
        prev = 0
        for h in hits:
            out.append(text[prev:h.start])
            out.append(f"<{h.category}>")
            prev = h.end
        out.append(text[prev:])
        return "".join(out)

    def check_outbound(self, parts: Dict[str, str]) -> DlpResult:
        """出站检查：对多段内容（query / context / output 等）统一扫描。

        **fail-closed 语义（目标 100% 拦截）**：
        - 命中任一规则 → block（`reason="hit"`）；
        - 引擎无任何规则（连内置保底都没有）且 fail_closed → block
          （`reason="no_rules"`，未配置 = 不安全，保守拒绝出站）。
        每检查一次落一条结构化审计记录（含命中规则/类别，不落明文敏感内容）。
        """
        all_hits: List[DlpHit] = []
        for part, text in (parts or {}).items():
            for h in self.inspect(text):
                # 带上所属 part，便于审计定位哪段内容触发拦截
                all_hits.append(DlpHit(
                    rule=h.rule, category=h.category, matched=h.matched,
                    start=h.start, end=h.end, part=part,
                ))
        all_hits.sort(key=lambda h: (h.start, h.end))

        if all_hits:
            decision, reason = "block", "hit"
        elif self.fail_closed and not self._rules:
            decision, reason = "block", "no_rules"
        else:
            decision, reason = "allow", "clean"

        audit: Dict[str, object] = {
            "audit_id": uuid.uuid4().hex,
            "ts": datetime.now(timezone.utc).isoformat(timespec="seconds"),
            "channel": "outbound",
            "template": self.template,
            "dlp_version": self.version,
            "decision": decision,
            "reason": reason,
            "parts": sorted(parts.keys()) if parts else [],
            "hit_count": len(all_hits),
            "hits": [h.to_dict() for h in all_hits],
        }
        self._audit_buffer.append(audit)
        if self.audit_hook is not None:
            self.audit_hook(audit)

        return DlpResult(decision=decision, reason=reason,
                         hits=all_hits, audit=audit)

    def drain_audit(self) -> List[Dict[str, object]]:
        """取走并清空已产生的审计记录（对接外部审计管道/日志）。"""
        buf, self._audit_buffer = self._audit_buffer, []
        return buf

    @property
    def rule_count(self) -> int:
        return len(self._rules)

    def rule_names(self) -> List[str]:
        return [r.name for r in self._rules]

    def __repr__(self) -> str:  # pragma: no cover - 调试辅助
        return (f"DlpEngine(template={self.template!r}, version={self.version!r}, "
                f"rules={len(self._rules)}, fail_closed={self.fail_closed})")


__all__ = [
    "DlpRuleKind",
    "DlpRule",
    "DlpHit",
    "DlpResult",
    "DlpEngine",
    "DLP_DEFAULT_RULES",
]