from __future__ import annotations

import json

from app.modules.agent.intent_router_v2.protocols import TextGenerator


class DocsLlmDisambiguator:
    _ALLOWED = {
        "SYSTEM_FLOW_EXPLAIN",
        "COMPONENT_EXPLAIN",
        "API_METHOD_EXPLAIN",
        "ENTITY_EXPLAIN",
        "RELATED_DOCS_EXPLAIN",
        "GENERIC_QA",
        "OPENAPI_METHOD_GENERATE",
        "OPENAPI_FRAGMENT_GENERATE",
    }

    def __init__(self, llm: TextGenerator) -> None:
        self._llm = llm

    def choose(self, payload: dict[str, object]) -> dict[str, str] | None:
        raw = self._llm.generate(
            "rag_docs_router_disambiguation_v1",
            json.dumps(payload, ensure_ascii=False),
            log_context="rag.intent_router_v2.disambiguate",
        ).strip()
        parsed = self._parse(raw)
        if parsed is None:
            return None
        return parsed

    def _parse(self, raw: str) -> dict[str, str] | None:
        candidate = self._strip_code_fence(raw)
        try:
            payload = json.loads(candidate)
        except json.JSONDecodeError:
            return None
        sub_intent = str(payload.get("sub_intent") or "").strip()
        if sub_intent not in self._ALLOWED:
            return None
        return {
            "sub_intent": sub_intent,
            "reason": str(payload.get("reason") or "").strip(),
            "confidence": str(payload.get("confidence") or "").strip(),
        }

    def _strip_code_fence(self, text: str) -> str:
        if not text.startswith("```"):
            return text
        lines = text.splitlines()
        if len(lines) >= 3 and lines[-1].strip() == "```":
            return "\n".join(lines[1:-1]).strip()
        return text
