"""对话层的工具注册表。

设计取舍（Why）：
- 工具 = 「JSON-Schema 声明 + 一个 Python 函数」，与 prodream dreami engine 的
  工具机制同形。LLM 只看到 schema，内核负责把调用路由到函数。
- 工具实现**不重写** analyze/diagnose/recheck 等能力，而是转调 main.py 里已有的
  endpoint 处理器——那四个能力是已验证资产（prompt 约束 + 数量校验 + 重试兜底
  全在里面），重写一份就是给自己留一个「改一条漏两条」的口子。
- 转调用的是延迟 import（函数内 `from main import ...`）：main.py 在模块顶层
  import 本模块，顶层再反向 import 会造成循环。延迟到调用时 main 已完全加载。
  **这个接缝同时保证了测试打桩点唯一**——test_demo.py 一律 monkeypatch
  `main.get_client`，工具走同一条路才不会绕过回归网。
- 工具给 LLM 的结果（summary）与给前端的结果（payload）分开：诊断 JSON 有 2.5k
  tokens，塞回对话历史会滚雪球；前端要渲染的原始结构又不能省。summary 是紧凑
  中文摘要（几十字），payload 原样回传。
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from fastapi import HTTPException

from schemas import (
    AnalyzeRequest,
    AnalyzeResponse,
    DiagnoseRequest,
    DiagnoseResponse,
    RecheckRequest,
    ReferenceRequest,
    ScaffoldRequest,
    TranslateRequest,
)


@dataclass
class ToolContext:
    """一次会话的业务上下文。原文由前端在建会话时给一次，
    之后所有工具从这里取——不让 LLM 在每轮参数里重抄整篇文书。"""

    prompt: str = ""
    word_limit: int | None = None
    paragraphs: list[str] = field(default_factory=list)
    constraints: list[str] = field(default_factory=list)
    rewrite_paragraphs: list[str] = field(default_factory=list)
    # 顾问确认/纠正后的语义锚点（与 paragraphs 等长，非空条目覆盖模型初判）
    confirmed_anchors: list[str] = field(default_factory=list)
    # 顾问分段要求：pid → 要求列表。优先级最高（PRD §10.1），必须透传到诊断与复检
    paragraph_constraints: dict[str, list[str]] = field(default_factory=dict)
    # 工具链累积的中间产物（dict 形态的 pydantic dump），供后续工具与前端复用
    analysis: dict[str, Any] | None = None
    diagnosis: dict[str, Any] | None = None

    def paragraph(self, paragraph_id: str) -> str:
        idx = _pid_index(paragraph_id)
        if idx is None or not (0 <= idx < len(self.paragraphs)):
            return ""
        return self.paragraphs[idx]

    def anchor(self, paragraph_id: str) -> str:
        """语义锚点：优先用理解阶段产出的，没有再退回顾问确认值。"""
        idx = _pid_index(paragraph_id)
        if idx is None:
            return ""
        brief = _brief(self.diagnosis, f"p{idx + 1}")
        if brief and str(brief.get("confirmed_meaning") or "").strip():
            return str(brief["confirmed_meaning"]).strip()
        if idx < len(self.paragraphs):
            paras = (self.analysis or {}).get("paragraphs") or []
            if idx < len(paras) and isinstance(paras[idx], dict):
                return str(paras[idx].get("semantic_anchor") or "").strip()
        return ""

    def goal(self, paragraph_id: str) -> str:
        brief = _brief(self.diagnosis, paragraph_id)
        return str((brief or {}).get("rewrite_goal") or "").strip()

    def paragraph_constraints_for(self, paragraph_id: str) -> list[str]:
        return list(self.paragraph_constraints.get(paragraph_id) or [])

    def anchors_for_diagnosis(self) -> list[str]:
        """诊断用的语义锚点：顾问确认值优先，没给就退回理解阶段的产出。

        退回这条不是兜底而是**保持既有行为**——前端原 requestDiagnosis() 就是
        `analysis.paragraphs.map(p => p.semantic_anchor || p.natural_meaning_zh)`，
        对话化改造不能让它变弱。"""
        if self.confirmed_anchors:
            return list(self.confirmed_anchors)
        out: list[str] = []
        for item in (self.analysis or {}).get("paragraphs") or []:
            if isinstance(item, dict):
                out.append(str(item.get("semantic_anchor") or item.get("natural_meaning_zh") or ""))
        return out


@dataclass
class ToolResult:
    ok: bool
    # 给 LLM 的紧凑摘要（进对话历史）
    summary: str
    # 给前端的原始结构（渲染诊断卡/理解卡；不进对话历史）
    payload: dict[str, Any] | None = None
    # 交互型工具：暂停本轮，等顾问选择后由下一条消息继续
    stop: bool = False
    clarify: dict[str, Any] | None = None
    error: str = ""


def _pid_index(paragraph_id: str) -> int | None:
    """p1 → 0。容忍模型给 'p1'/'P1'/'1'，给不出就返回 None（不猜）。"""
    token = str(paragraph_id or "").strip().lower().lstrip("p")
    return int(token) - 1 if token.isdigit() else None


def _brief(diagnosis: dict[str, Any] | None, paragraph_id: str) -> dict[str, Any] | None:
    for brief in (diagnosis or {}).get("paragraph_briefs") or []:
        if isinstance(brief, dict) and brief.get("paragraph_id") == paragraph_id:
            return brief
    return None


def _call(handler: Any, request: Any) -> dict[str, Any]:
    """转调 endpoint 处理器，把 HTTP 语义的错误翻成工具错误。

    工具失败必须让 LLM 看见（SKILL.md 诚实性硬约束第 4 条：不许用自己
    的想法把结果补齐），所以这里返回 error 而不是抛异常中断整轮对话。"""
    try:
        resp = handler(request)
    except HTTPException as exc:
        return {"error": str(exc.detail)}
    except Exception as exc:  # 未预期：同样转成可读错误，不让它打断会话
        return {"error": f"{exc.__class__.__name__}：{exc}"}
    return {"data": resp.model_dump() if hasattr(resp, "model_dump") else resp}


# --------------------------------------------------------------------------
# 工具实现
def _analyze(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import analyze  # 延迟 import：见模块 docstring

    if not ctx.paragraphs:
        return ToolResult(ok=False, error="还没有拿到文书原文，无法分析", summary="")
    out = _call(
        analyze,
        AnalyzeRequest(
            prompt=ctx.prompt,
            word_limit=ctx.word_limit,
            paragraphs=ctx.paragraphs,
            constraints=ctx.constraints,
        ),
    )
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    ctx.analysis = data
    paras = data.get("paragraphs") or []
    summary = (
        f"理解完成：{len(paras)} 段，摘要「{str(data.get('essay_summary') or '')[:80]}」；"
        f"各段语义锚点：" + "；".join(
            f"p{i + 1}={str((p or {}).get('semantic_anchor') or '')[:40]}"
            for i, p in enumerate(paras)
        )
    )
    return ToolResult(ok=True, summary=summary, payload={"kind": "analysis", "data": data})


def _diagnose(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import diagnose

    if not ctx.paragraphs:
        return ToolResult(ok=False, error="还没有拿到文书原文，无法诊断", summary="")
    initial = None
    if ctx.analysis:
        try:
            initial = AnalyzeResponse.model_validate(ctx.analysis)
        except Exception:
            initial = None  # 结构不符就当作没有理解结果，不阻断诊断
    out = _call(
        diagnose,
        DiagnoseRequest(
            prompt=ctx.prompt,
            word_limit=ctx.word_limit,
            paragraphs=ctx.paragraphs,
            confirmed_anchors=ctx.anchors_for_diagnosis(),
            global_constraints=ctx.constraints,
            paragraph_constraints=ctx.paragraph_constraints,
            initial_analysis=initial,
        ),
    )
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    ctx.diagnosis = data
    patterns = data.get("patterns") or []
    summary = (
        f"诊断完成：{len(patterns)} 个 Pattern。"
        + "；".join(
            f"{p.get('pattern_id')} {p.get('name')}"
            f"（段 {'/'.join(p.get('affected_paragraphs') or []) or '未标'}）："
            f"{str(p.get('transformation_rule') or '')[:40]}"
            for p in patterns
        )
        + f"。总述：{str(data.get('overall_diagnosis') or '')[:120]}"
    )
    return ToolResult(ok=True, summary=summary, payload={"kind": "diagnosis", "data": data})


def _scaffold(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import scaffold

    pid = str(args.get("paragraph_id") or "").strip()
    text = ctx.paragraph(pid)
    if not text:
        return ToolResult(ok=False, error=f"找不到段落 {pid or '(未给段号)'}", summary="")
    out = _call(
        scaffold,
        ScaffoldRequest(
            paragraph_id=f"p{_pid_index(pid) + 1}",
            original_text=text,
            semantic_anchor=ctx.anchor(pid),
            rewrite_goal=str(args.get("rewrite_goal") or ctx.goal(pid)),
            global_constraints=ctx.constraints,
            paragraph_constraints=ctx.paragraph_constraints_for(f"p{_pid_index(pid) + 1}"),
        ),
    )
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    return ToolResult(
        ok=True,
        summary=f"{pid} 写作起点：{str(data.get('scaffold') or '')[:200]}",
        payload={"kind": "scaffold", "data": data},
    )


def _reference(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import reference

    pid = str(args.get("paragraph_id") or "").strip()
    text = ctx.paragraph(pid)
    if not text:
        return ToolResult(ok=False, error=f"找不到段落 {pid or '(未给段号)'}", summary="")
    out = _call(
        reference,
        ReferenceRequest(
            paragraph_id=f"p{_pid_index(pid) + 1}",
            original_text=text,
            semantic_anchor=ctx.anchor(pid),
            rewrite_goal=str(args.get("rewrite_goal") or ctx.goal(pid)),
            global_constraints=ctx.constraints,
            paragraph_constraints=ctx.paragraph_constraints_for(f"p{_pid_index(pid) + 1}"),
        ),
    )
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    return ToolResult(
        ok=True,
        summary=f"{pid} 参考片段：{str(data.get('reference_snippet') or '')[:200]}",
        payload={"kind": "reference", "data": data},
    )


def _translate(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import translate

    pid = str(args.get("paragraph_id") or "").strip()
    text = str(args.get("text") or "").strip()
    if pid and not text:
        # 优先原文；顾问已在工作台改写则退回改写稿（他想核对的是自己改的那版）
        text = ctx.paragraph(pid)
        idx = _pid_index(pid)
        if not text and idx is not None and idx < len(ctx.rewrite_paragraphs):
            text = ctx.rewrite_paragraphs[idx]
    if not text:
        return ToolResult(ok=False, error="没有可翻译的文本", summary="")
    out = _call(translate, TranslateRequest(paragraph_id=pid, text=text))
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    return ToolResult(
        ok=True,
        summary=f"中文对照：{str(data.get('translation') or '')[:300]}",
        payload={"kind": "translation", "data": data},
    )


def _recheck(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    from main import recheck

    if not ctx.rewrite_paragraphs:
        return ToolResult(
            ok=False,
            error="还没有拿到改写稿。复检需要顾问先在工作台完成各段改写",
            summary="",
        )
    if len(ctx.rewrite_paragraphs) != len(ctx.paragraphs):
        return ToolResult(
            ok=False,
            error=f"原文 {len(ctx.paragraphs)} 段、改写 {len(ctx.rewrite_paragraphs)} 段，数量不一致，无法复检",
            summary="",
        )
    initial = None
    if ctx.diagnosis:
        try:
            initial = DiagnoseResponse.model_validate(ctx.diagnosis)
        except Exception:
            initial = None  # 少了诊断不影响复检其余六个维度
    out = _call(
        recheck,
        RecheckRequest(
            prompt=ctx.prompt,
            word_limit=ctx.word_limit,
            original_paragraphs=ctx.paragraphs,
            rewrite_paragraphs=ctx.rewrite_paragraphs,
            confirmed_anchors=ctx.anchors_for_diagnosis(),
            global_constraints=ctx.constraints,
            paragraph_constraints=ctx.paragraph_constraints,
            # 不带 diagnosis 则「原 Pattern 是否缓解」这一维失去对照，复检会变弱
            diagnosis=initial,
        ),
    )
    if "error" in out:
        return ToolResult(ok=False, error=out["error"], summary="")
    data = out["data"]
    checks = data.get("global_checks") or {}
    targets = data.get("revision_targets") or []
    failed = [k for k, v in checks.items() if v == "fail"]
    summary = (
        f"复检结果：{data.get('status')}"
        f"（未通过维度：{'/'.join(failed) or '无'}）。"
        + (
            "返工目标：" + "；".join(
                f"{t.get('paragraph_id')} {str(t.get('single_revision_goal') or '')[:60]}"
                for t in targets
            )
            if targets
            else "无需返工"
        )
    )
    return ToolResult(ok=True, summary=summary, payload={"kind": "recheck", "data": data})


def _record_note(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    """把顾问口头提的要求落成结构化约束。

    存在的理由：顾问说「第三段理解不对，重点不是接受不确定」时，如果只是口头
    应下来，这条要求到不了诊断——prompt 里顾问约束优先级最高（PRD §10.1），
    丢了它等于把顾问的纠正当耳边风。落进 ctx 后由 _diagnose/_recheck 透传，
    同时回传 payload 让前端同步自己的副本（下一轮前端会重发上下文，
    两边不一致就会互相覆盖）。"""
    note = str(args.get("note") or "").strip()
    if not note:
        return ToolResult(ok=False, error="note 不能为空", summary="")
    pid = str(args.get("paragraph_id") or "").strip()
    if pid:
        idx = _pid_index(pid)
        # 段号给错时**报错而不是降级成全篇要求**：模型幻觉出一个不存在的段号，
        # 静默接受会让一条本该只作用于 p3 的要求污染全文
        if idx is None or not (0 <= idx < len(ctx.paragraphs)):
            return ToolResult(
                ok=False,
                error=f"找不到段落 {pid}（全文共 {len(ctx.paragraphs)} 段），请用 p1/p2 这样的段号；全篇要求请留空",
                summary="",
            )
        key = f"p{idx + 1}"
        bucket = ctx.paragraph_constraints.setdefault(key, [])
        if note not in bucket:
            bucket.append(note)
        return ToolResult(
            ok=True,
            summary=f"已记入 {key} 的要求：{note}",
            payload={"kind": "constraint", "data": {"paragraph_id": key, "note": note}},
        )
    if note not in ctx.constraints:
        ctx.constraints.append(note)
    data = {"paragraph_id": "", "note": note}
    summary = f"已记入全局要求：{note}"
    return ToolResult(ok=True, summary=summary, payload={"kind": "constraint", "data": data})


def _confirm(ctx: ToolContext, args: dict[str, Any]) -> ToolResult:
    """交互型工具：产出选项按钮并暂停本轮（对应 prodream 的 stopAfterToolCall）。"""
    question = str(args.get("question") or "").strip()
    choices = [str(c).strip() for c in (args.get("choices") or []) if str(c).strip()]
    if not question or not choices:
        return ToolResult(ok=False, error="let_user_confirm 需要 question 与非空 choices", summary="")
    return ToolResult(
        ok=True,
        summary=f"已向顾问提问：{question}",
        payload={"kind": "confirm", "data": {"question": question, "choices": choices}},
        stop=True,
        clarify={"question": question, "choices": choices},
    )


# --------------------------------------------------------------------------
# 注册表
@dataclass(frozen=True)
class ToolSpec:
    name: str
    description: str
    parameters: dict[str, Any]
    run: Any


TOOLS: list[ToolSpec] = [
    ToolSpec(
        name="analyze_essay",
        description=(
            "读懂全文：给出各段的自然中文理解与「必须保留的核心意思」（语义锚点），"
            "并判断文章如何回应 Essay Prompt。不产出 AI 味问题、不做质量评价。"
            "顾问首次要求「分析/看看这篇」时先调它。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "focus": {"type": "string", "description": "顾问本次特别关心的点，可留空"},
            },
        },
        run=_analyze,
    ),
    ToolSpec(
        name="diagnose_essay",
        description=(
            "诊断 AI 味：找出文章中实际命中的生成式写作 Pattern（如特质宣告、完美成长弧、"
            "排比对称、格言体），每个给出逐字原文证据、为什么像 AI 写、为什么影响本人感、"
            "以及对应的人工改写动作；同时给出每段的改写目标与批注。"
            "顾问问「哪里像 AI / 有什么问题」时调它。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "focus": {"type": "string", "description": "顾问本次特别关心的点，可留空"},
            },
        },
        run=_diagnose,
    ),
    ToolSpec(
        name="get_writing_scaffold",
        description=(
            "给某一段一个轻量写作起点（填空/思考顺序/一句不完整起笔）。"
            "顾问说「这段不知道怎么改/给个方向」时调它。不产出完整段落。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "paragraph_id": {"type": "string", "description": "段落 id，如 p2"},
                "rewrite_goal": {"type": "string", "description": "本轮这一段要完成的变化，可留空（留空则用诊断给出的目标）"},
            },
            "required": ["paragraph_id"],
        },
        run=_scaffold,
    ),
    ToolSpec(
        name="get_reference_snippet",
        description=(
            "给某一段 1–2 句局部参考，用来说明改写动作长什么样。"
            "顾问卡住、说「给个例子/参考」时调它。不是整段答案。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "paragraph_id": {"type": "string", "description": "段落 id，如 p2"},
                "rewrite_goal": {"type": "string", "description": "本轮这一段要完成的变化，可留空"},
            },
            "required": ["paragraph_id"],
        },
        run=_reference,
    ),
    ToolSpec(
        name="translate_to_chinese",
        description=(
            "把原文某段（或顾问指定的英文文本）译成中文，用于核对改写后有没有偏离原意。"
            "忠实直译，不做评价与润色。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "paragraph_id": {"type": "string", "description": "段落 id，如 p2；与 text 二选一"},
                "text": {"type": "string", "description": "要翻译的英文文本；留空则翻译 paragraph_id 对应的原文"},
            },
        },
        run=_translate,
    ),
    ToolSpec(
        name="recheck_rewrite",
        description=(
            "对顾问已完成的整篇改写稿做全文复检：语义保留、AI 味缓解、有无新替代模板、"
            "全篇声音一致性、是否套用支架/参考、字数。只在顾问说「改完了/提交复检」时调它。"
            "注意：你无法替顾问改稿，此工具读的是顾问在工作台的改写稿。"
        ),
        parameters={"type": "object", "properties": {}},
        run=_recheck,
    ),
    ToolSpec(
        name="record_advisor_note",
        description=(
            "把顾问提出的一条要求/纠正记入工作记录，后续诊断与复检都会遵守。"
            "顾问说「这段理解不对」「保留某个表达」「语言更口语化」这类话时必须调用它——"
            "只口头答应不记录，要求会丢失。"
            "针对具体段落就传 paragraph_id，全篇性的就留空。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "note": {"type": "string", "description": "顾问的要求原文或其准确转述，一句话"},
                "paragraph_id": {"type": "string", "description": "只针对某一段时填，如 p2；全篇要求留空"},
            },
            "required": ["note"],
        },
        run=_record_note,
    ),
    ToolSpec(
        name="let_user_confirm",
        description=(
            "向顾问提问并给出可点选的选项，然后暂停等你回答。"
            "用于需要顾问拍板的场合：选哪一段、确认改写方向、下一步做什么。"
            "调用后本轮结束，顾问的选择会作为下一条消息回到你这里。"
        ),
        parameters={
            "type": "object",
            "properties": {
                "question": {"type": "string", "description": "要问顾问的问题，一句话"},
                "choices": {
                    "type": "array",
                    "items": {"type": "string"},
                    "description": "2–4 个选项，每个是一句可直接执行的话",
                },
            },
            "required": ["question", "choices"],
        },
        run=_confirm,
    ),
]

_BY_NAME = {spec.name: spec for spec in TOOLS}


def openai_tools() -> list[dict[str, Any]]:
    """转成 OpenAI 兼容的 tools 声明（engine 侧剥掉的 execUrl 在这里由注册表承担）。"""
    return [
        {
            "type": "function",
            "function": {
                "name": spec.name,
                "description": spec.description,
                "parameters": spec.parameters,
            },
        }
        for spec in TOOLS
    ]


def execute(name: str, args: dict[str, Any], ctx: ToolContext) -> ToolResult:
    spec = _BY_NAME.get(name)
    if spec is None:
        return ToolResult(ok=False, error=f"未知工具 {name}", summary="")
    try:
        return spec.run(ctx, args or {})
    except Exception as exc:  # 兜底：单个工具崩溃不能带走整轮对话
        return ToolResult(ok=False, error=f"{exc.__class__.__name__}：{exc}", summary="")