"""阶段 3：逐图框审查（LLM + 受控只读取证子循环）。

流程：
  1. 用清单生成系统提示，喂入结构化证据 + 截图。
  2. LLM 若请求只读取证工具 → 执行后把观测回灌，再问（有界循环）。
  3. LLM 不再请求工具时，其输出（受约束的 JSON）即为该图框的审查结论。

结论是**结构化对象** `{"summary": str, "issues": [dict], "raw": str}`，
不再依赖第二次 LLM 调用来抽取结构（见 `findings.parse_conclusion`）。

审查在**绘图区已最大化**的前提下进行（见 `evidence.maximized_drawing_area`），
子循环里的放大截图因此是高清的；若 LLM 动过视口，结束时复位到该图框整幅视图。
"""

import json

from .. import config, evidence_model, inspect
from ..cad_tools import READ_ONLY_TOOLS, is_write_tool
from ..findings import parse_conclusion
from ..llm_client import LlmClient, downsample_image_b64
from ..mcp_client import McpError, McpToolError, result_images, result_texts

# 单次审查回灌给 LLM 的图片数上限（控制视觉 token 成本）
_MAX_OBS_IMAGES = config.REVIEW_OBS_IMAGES

# 会改动视口的只读取证工具：用过之后须把视口复位，别把用户的 CAD 留在放大处
_VIEW_TOOLS = frozenset({"zoom_to_window", "zoom_to_entity", inspect.TOOL_NAME})

# 启用 inspect_region 时，这三个裸视口工具不再直接提供给 LLM（由 inspect_region 统一代理），
# 但仍留在只读白名单里，作为兜底仍然可被调用。
_RAW_VIEW_TOOLS = frozenset(
    {"zoom_to_window", "zoom_to_entity", "get_viewport_screenshot"}
)


def evidence_payload(evidence):
    """把证据整理成给 LLM 的 JSON（剔除 base64 图片，避免撑爆上下文）。"""
    frame = evidence.get("frame") or {}
    return {
        "source_type": evidence.get("source_type"),
        "frame_index": evidence.get("index"),
        "frame": {
            k: frame.get(k)
            for k in ("sheet_size", "kind", "layer", "bbox", "confidence")
        },
        "entities": evidence.get("entities"),
        "texts": evidence.get("texts"),
        "dimensions": evidence.get("dimensions"),
        "blocks": evidence.get("blocks"),
        "polylines": evidence.get("polylines"),
    }


def _run_readonly_tool(ctx, name, args, image_budget, *, evidence=None, state=None):
    """执行一个只读取证工具，返回回灌给 LLM 的 content 列表。

    `inspect_region` 是 **agent 侧虚拟工具**（不走 MCP），在只读白名单校验**之前**拦截：
    它内部只调白名单内的只读 MCP 工具。

    离线源（PDF / 图片，`ctx.tools is None`）没有任何 MCP 工具可用 → 一律拒绝；
    `inspect_region` 依赖 CAD 截图，同样不可用。
    """
    if ctx.tools is None:
        return [{"type": "text", "text": f"当前为离线源，工具 {name} 不可用。"}], image_budget

    if name == inspect.TOOL_NAME:
        if not config.INSPECT_ENABLED or evidence is None or state is None:
            return [{"type": "text", "text": "inspect_region 当前不可用。"}], image_budget
        return inspect.inspect_region(ctx, evidence, args, image_budget, state)

    if name not in READ_ONLY_TOOLS or is_write_tool(name):
        return [{"type": "text", "text": f"工具 {name} 不在只读取证白名单内，已拒绝。"}], image_budget

    try:
        resp = ctx.tools.call(name, args)
    except (McpError, McpToolError) as exc:
        return [{"type": "text", "text": f"工具 {name} 执行失败: {exc}"}], image_budget

    content = []
    for text in result_texts(resp):
        content.append({"type": "text", "text": text})
    for b64, mime in result_images(resp):
        if image_budget[0] <= 0:
            content.append(
                {
                    "type": "text",
                    "text": "（截图预算已用完，后续截图已省略；请基于已有信息给出结论，不要再请求截图。）",
                }
            )
            break
        image_budget[0] -= 1
        nb, nm = downsample_image_b64(b64, mime)
        content.append(LlmClient.image_content(nb, nm))
    if not content:
        content.append({"type": "text", "text": "（无返回内容）"})
    return content, image_budget


def _review_tools(ctx):
    """审查子循环提供给 LLM 的工具 schema。

    启用 `inspect_region` 时，用它**取代**三个裸视口工具（`zoom_to_window` /
    `zoom_to_entity` / `get_viewport_screenshot`）——一个语义化入口比三个裸工具更好用，
    且可读性由 agent 保证；关闭时维持原样（退回旧行为）。

    离线源（`ctx.tools is None`）不提供任何工具（也没有 CAD 可 zoom / 截图）。
    """
    if ctx.tools is None:
        return []
    if config.INSPECT_ENABLED:
        names = READ_ONLY_TOOLS - _RAW_VIEW_TOOLS
        return ctx.tools.openai_tools(names=names) + [inspect.TOOL_SCHEMA]
    return ctx.tools.openai_tools(names=READ_ONLY_TOOLS)


def _reset_viewport(ctx, evidence):
    """把视口复位到该图框的整幅视图（审查中的放大不应残留给用户）。"""
    win = evidence.get("window")
    if not win or ctx.tools is None:
        return
    try:
        ctx.tools.call("zoom_to_window", dict(win, margin=0.02))
    except (McpError, McpToolError):
        pass


def review_frame(ctx, evidence, system_prompt, log=print):
    """审查单个图框，返回结构化结论 `{"summary", "issues", "raw"}`。

    审查中的放大取证会改动视口，故在 finally 里复位——**仅当确实动过视口**
    （LLM 只查结构化数据时不动用户视图）。
    """
    view_changed = [False]
    try:
        return _review_frame(ctx, evidence, system_prompt, log, view_changed)
    finally:
        if view_changed[0]:
            _reset_viewport(ctx, evidence)


def _review_frame(ctx, evidence, system_prompt, log, view_changed):
    payload = evidence_payload(evidence)
    user_content = [
        {
            "type": "text",
            "text": "以下是该图框的审查输入数据(JSON)：\n"
            + json.dumps(payload, ensure_ascii=False, indent=2),
        }
    ]
    for img in evidence_model.evidence_images(evidence):
        nb, nm = downsample_image_b64(img["b64"], img.get("mime") or "image/png")
        label = img.get("label")
        if label:
            user_content.append({"type": "text", "text": f"（截图：{label}）"})
        user_content.append(LlmClient.image_content(nb, nm))

    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]

    tools_schema = _review_tools(ctx)
    image_budget = [_MAX_OBS_IMAGES]
    state = inspect.InspectState(evidence)

    for _ in range(max(0, ctx.max_iters)):
        resp = ctx.llm.chat(messages, tools=tools_schema)
        msg = resp.choices[0].message
        messages.append(msg)

        # 仅在「不再请求工具」的回合解析结论；中间工具回合的 content 永不解析
        if not msg.tool_calls:
            return parse_conclusion(msg.content or "", log=log)

        for tool_call in msg.tool_calls:
            name = tool_call.function.name
            try:
                args = json.loads(tool_call.function.arguments or "{}")
            except json.JSONDecodeError:
                args = {}
            log(f"      [取证] {name} {args if args else ''}")
            if name in _VIEW_TOOLS:
                view_changed[0] = True
            content, image_budget = _run_readonly_tool(
                ctx, name, args, image_budget, evidence=evidence, state=state
            )
            messages.append(
                {"role": "tool", "tool_call_id": tool_call.id, "content": content}
            )

    # 达到回合上限：强制收口（不传 tools，保证不再请求工具，输出必为终稿 JSON）
    messages.append(
        {
            "role": "user",
            "content": "已达到取证回合上限，请基于现有信息直接给出该图框的审查结论（JSON）。",
        }
    )
    resp = ctx.llm.chat(messages)
    return parse_conclusion(resp.choices[0].message.content or "", log=log)