"""LLM 兜底打 role：对低置信段落做二次确认。"""

from __future__ import annotations
import logging
from fastapi_modules.fastapi_leaudit.govdoc_engine.models import Document, Role
from fastapi_modules.fastapi_leaudit.govdoc_engine.llm.client import LlmClient, _format_exc

_log = logging.getLogger(__name__)


VALID_ROLES = [
    "title", "doc_number", "recipient",
    "heading_1", "heading_2", "heading_3", "heading_4",
    "body", "attachment_marker", "signature", "date",
    "no_text_marker", "unknown",
]


_PROMPT = """你是公文格式专家。下面是一份公文的段落列表，请为指定的"待定段落"判断其角色。

【全文段落（带索引和当前规则推测）】
{context}

【待定段落 idx={idx}】
文本: {text}
当前推测角色: {current_role}（置信度 {conf:.2f}）

【角色取值范围】
{roles}

请综合公文结构判断该段落最可能的角色。

仅以 JSON 输出：
{{"role": "<角色>", "confidence": <0-1 浮点数>, "reason": "<简短理由>"}}
"""


class LlmTagger:
    def __init__(self, client: LlmClient):
        self.client = client

    def _build_prompt(self, doc: Document, target_idx: int) -> tuple[str, "object"]:
        ctx_lines = []
        for p in doc.paragraphs:
            tag = "← 待定" if p.index == target_idx else ""
            ctx_lines.append(f"[{p.index}] role={p.role} text={p.text[:60]} {tag}")
        ctx = "\n".join(ctx_lines)
        target = doc.paragraphs[target_idx]
        prompt = _PROMPT.format(
            context=ctx,
            idx=target_idx,
            text=target.text,
            current_role=target.role or "unknown",
            conf=target.role_confidence,
            roles=", ".join(VALID_ROLES),
        )
        return prompt, target

    def _interpret(self, resp: dict, target) -> tuple[Role, float]:
        role = resp.get("role", "unknown")
        if role not in VALID_ROLES:
            role = "unknown"
        conf = float(resp.get("confidence", 0.5))
        return role, conf  # type: ignore[return-value]

    def disambiguate(self, doc: Document, target_idx: int) -> tuple[Role, float]:
        prompt, target = self._build_prompt(doc, target_idx)
        label = f"role_tag_p{target_idx}"
        try:
            resp = self.client.chat_json(
                [{"role": "user", "content": prompt}], label=label,
            )
        except Exception as e:
            _log.warning("Role disambiguation skipped (LLM error): %s", _format_exc(e))
            return target.role or "unknown", target.role_confidence  # type: ignore[return-value]
        return self._interpret(resp, target)

    async def disambiguate_async(
        self, doc: Document, target_idx: int
    ) -> tuple[Role, float]:
        prompt, target = self._build_prompt(doc, target_idx)
        label = f"role_tag_p{target_idx}"
        try:
            resp = await self.client.chat_json_async(
                [{"role": "user", "content": prompt}], label=label,
            )
        except Exception as e:
            _log.warning("Role disambiguation skipped (LLM error): %s", _format_exc(e))
            return target.role or "unknown", target.role_confidence  # type: ignore[return-value]
        return self._interpret(resp, target)
