# -*- coding: utf-8 -*-
"""
tool_bridge.py — Pi Agent 兜底能力 P0 PoC：工具桥与 callId 凭证
================================================================

对应 `docs/architecture/fallback.md` 与 `P1-DESIGN.md`：

- §4.3 工具桥暴露面：定义围墙内工具白名单注册表（名称 / 权力等级 P0-P2 /
  参数 schema / 是否需确认）。产品态里这些工具经由 HTTP 桥暴露给 Pi；
  PoC 阶段由 mock executor 直接调用本模块，接口签名保持一致。
- §4.4/§4.5 证据链：每次工具调用签发 callId 凭证，落 runs/<runId>/calls.jsonl
  （actor=pi-fallback:<runId>，入参摘要、结果摘要）。Pi 报告中引用的 callId
  必须真实存在 —— 「无凭证的成果声明 → 判 invalid」（GOAL.md 越狱测试之三：
  伪造成果无 callId）。

callId 设计说明：真实 pi 事件流自带 toolCallId（Agent-A 侦察 §5-7）；PoC 的
桥层凭证独立签发（callId 前缀 "call-"），两者在 calls.jsonl 里可互相登记关联，
防止「绕过桥直接向 LLM 要一个编造的数字」。
"""

from __future__ import annotations

import hashlib
import json
import re
import time
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

# ---------------------------------------------------------------------------
# 工具白名单注册表（§4.3 暴露面；权力等级沿用 _POWER_MAP 语义 P0 只读 < P1 计划 < P2 执行）
# ---------------------------------------------------------------------------

@dataclass(frozen=True)
class ToolSpec:
    """单个围墙内工具的登记项。"""
    name: str
    power: str                 # "P0" 只读 / "P1" 产出草稿 / "P2" 执行写
    needs_confirm: bool        # P2 级工具必须 True（确认卡机制，方案 §4.4）
    params_schema: dict        # JSON-schema 风格的参数说明（PoC 只做存在性登记）
    description: str = ""


TOOL_REGISTRY: dict[str, ToolSpec] = {t.name: t for t in [
    ToolSpec(
        name="aps_query", power="P0", needs_confirm=False,
        params_schema={"sql_like": "string"},
        description="世界状态只读视图（包装 master.query / data.analyze 只读集）",
    ),
    ToolSpec(
        name="knowledge_query", power="P0", needs_confirm=False,
        params_schema={"q": "string"},
        description="知识库查询（包装 /api/rag/query，产品态需 ragScopes 鉴权）",
    ),
    ToolSpec(
        name="fs_read", power="P0", needs_confirm=False,
        params_schema={"path": "string(限 run 目录内)"},
        description="沙箱 L2 三区只读（inbox 注入数据主要入口）",
    ),
    ToolSpec(
        name="fs_write", power="P1", needs_confirm=False,
        params_schema={"path": "string(限 work/ 或 outbox/)", "content": "string"},
        description="沙箱内可写工作区/产物出口写文件；越界即拒绝",
    ),
    ToolSpec(
        name="checkpoint_create", power="P1", needs_confirm=False,
        params_schema={},
        description="建回滚锚点，Pi 可主动调用（产品态对应 state checkpoints）",
    ),
    ToolSpec(
        name="shell_run", power="P2", needs_confirm=True,
        params_schema={"cmd": "string(须命中命令白名单正则)"},
        description="白名单命令执行（python 脚本、xlsx 处理等显式登记命令）；本 PoC 默认全禁",
    ),
    ToolSpec(
        name="aps_invoke", power="P2", needs_confirm=True,
        params_schema={"intent": "string", "params": "object"},
        description="Pi 写世界的唯一方式（包装 /api/agent/invoke，内部仍过 _POWER_MAP/确认卡）",
    ),
    ToolSpec(
        name="report_emit", power="P1", needs_confirm=False,
        params_schema={"md": "string"},
        description="产物唯一出口：报告强制走验证器（§4.3 report_emit / §4.5 REPORT 审计）",
    ),
]}


def check_tool_registered(tool_name: str) -> ToolSpec:
    """校验工具在白名单注册表内；未登记即拒绝（L1 工具层围墙的桥侧执行点）。"""
    spec = TOOL_REGISTRY.get(tool_name)
    if spec is None:
        raise ToolBridgeViolation(f"工具未在白名单注册表登记: {tool_name}")
    return spec


class ToolBridgeViolation(Exception):
    """工具桥违规（未登记工具 / 凭证缺失 / 伪造成果）。一律按失败处理。"""


# ---------------------------------------------------------------------------
# callId 凭证
# ---------------------------------------------------------------------------

_CALLS_FILE = "calls.jsonl"
# 报告里引用 callId 的约定格式： [callId: call-xxxx] 或 （凭证 call-xxxx）
_CALLID_RE = re.compile(r"call-[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")


@dataclass
class ToolCall:
    """一次（模拟）工具调用的凭证记录。"""
    call_id: str
    tool: str
    run_id: str
    ts: float
    params_digest: str = ""     # 入参摘要（sha256 前 16 位，不落明文防泄露敏感参数）
    status: str = "issued"      # issued -> completed / failed
    result_digest: str = ""     # 结果摘要


class ToolBridge:
    """桥本体：签发 callId、落 calls.jsonl、校验报告引用。"""

    def __init__(self, run_id: str, run_dir: Path):
        self.run_id = run_id
        self.calls_path = Path(run_dir) / _CALLS_FILE

    # -- 签发与完成 --------------------------------------------------------

    def issue_call(self, tool_name: str, params: Optional[dict] = None) -> str:
        """
        每次工具调用前签发 callId 凭证并落账（status=issued）。
        未登记工具抛 ToolBridgeViolation。
        返回 callId（uuid4，"call-" 前缀）。
        """
        spec = check_tool_registered(tool_name)
        call_id = "call-" + str(uuid.uuid4())
        rec = ToolCall(
            call_id=call_id, tool=spec.name, run_id=self.run_id, ts=time.time(),
            params_digest=self._digest(params),
        )
        self._append(rec.__dict__)
        return call_id

    def complete_call(self, call_id: str, result, ok: bool = True) -> None:
        """工具调用完成后补记结果摘要（status=completed/failed，追加一条记录）。"""
        known = {c["call_id"] for c in self.list_calls()}
        if call_id not in known:
            raise ToolBridgeViolation(f"complete_call: 未知 callId {call_id}")
        rec = {
            "call_id": call_id, "run_id": self.run_id, "ts": time.time(),
            "status": "completed" if ok else "failed",
            "result_digest": self._digest(result),
        }
        self._append(rec)

    # -- 查询与校验 --------------------------------------------------------

    def list_calls(self) -> list:
        """读出本次运行的全部凭证记录（calls.jsonl 逐行解析）。"""
        if not self.calls_path.exists():
            return []
        out = []
        with open(self.calls_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if line:
                    out.append(json.loads(line))
        return out

    def validate_report_citations(self, report_text: str) -> dict:
        """
        成果校验：Pi 报告中引用的每个 callId 都必须真实存在于 calls.jsonl。

        返回 {"valid": bool, "cited": [...], "missing": [...], "issued": [...]}。
        规则（GOAL.md 越狱测试判据）：
          - 引用不存在 → missing 非空 → valid=False（伪造成果）；
          - 报告声明了工具成果却一个 callId 都不引用 → 由调用方结合业务判定，
            本函数给出 cited 列表供其检查。
        """
        cited = sorted(set(_CALLID_RE.findall(report_text)))
        issued = {c["call_id"] for c in self.list_calls()}
        missing = [c for c in cited if c not in issued]
        return {
            "valid": not missing,
            "cited": cited,
            "missing": missing,
            "issued": sorted(issued),
        }

    # -- 内部 --------------------------------------------------------------

    @staticmethod
    def _digest(obj) -> str:
        try:
            blob = json.dumps(obj, ensure_ascii=False, sort_keys=True, default=str)
        except Exception:
            blob = str(obj)
        return hashlib.sha256(blob.encode("utf-8")).hexdigest()[:16]

    def _append(self, rec: dict) -> None:
        with open(self.calls_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(rec, ensure_ascii=False) + "\n")