from __future__ import annotations

from app.modules.agent.runtime.docs_qa_pipeline.models import DocsEvidenceBundle


class DocsEvidenceBuilder:
    _CAPS = {
        "SYSTEM_FLOW_EXPLAIN": {"documents": 1, "facts": 1, "entities": 0, "workflows": 2, "relations": 2, "chunks": 2},
        "COMPONENT_EXPLAIN": {"documents": 1, "facts": 3, "entities": 0, "workflows": 0, "relations": 2, "chunks": 2},
        "API_METHOD_EXPLAIN": {"documents": 1, "facts": 3, "entities": 0, "workflows": 1, "relations": 0, "chunks": 2},
        "ENTITY_EXPLAIN": {"documents": 1, "facts": 1, "entities": 2, "workflows": 0, "relations": 2, "chunks": 1},
        "RELATED_DOCS_EXPLAIN": {"documents": 2, "facts": 0, "entities": 1, "workflows": 0, "relations": 4, "chunks": 1},
        "GENERIC_QA": {"documents": 2, "facts": 0, "entities": 0, "workflows": 0, "relations": 0, "chunks": 1},
        "OPENAPI_METHOD_GENERATE": {"documents": 1, "facts": 4, "entities": 0, "workflows": 0, "relations": 0, "chunks": 2},
        "OPENAPI_FRAGMENT_GENERATE": {"documents": 1, "facts": 4, "entities": 0, "workflows": 0, "relations": 0, "chunks": 2},
    }
    _LAYER_TO_BUCKET = {
        "D1_DOCUMENT_CATALOG": "documents",
        "D2_FACT_INDEX": "facts",
        "D3_ENTITY_CATALOG": "entities",
        "D4_WORKFLOW_INDEX": "workflows",
        "D5_RELATION_GRAPH": "relations",
        "D0_DOC_CHUNKS": "chunks",
    }

    def build(self, *, intent: str, sub_intent: str, raw_rows: list[dict]) -> DocsEvidenceBundle:
        buckets = {name: [] for name in self._LAYER_TO_BUCKET.values()}
        for row in raw_rows:
            bucket = self._LAYER_TO_BUCKET.get(str(row.get("layer") or ""))
            if bucket is not None:
                buckets[bucket].append(self._normalize_row(row))
        support_paths = self._support_paths(sub_intent, buckets)
        caps = self._CAPS.get(sub_intent, self._CAPS["GENERIC_QA"])
        return DocsEvidenceBundle(
            intent=intent,
            sub_intent=sub_intent,
            documents=buckets["documents"][: caps["documents"]],
            facts=buckets["facts"][: caps["facts"]],
            entities=buckets["entities"][: caps["entities"]],
            workflows=buckets["workflows"][: caps["workflows"]],
            relations=buckets["relations"][: caps["relations"]],
            chunks=self._select_chunks(buckets["chunks"], support_paths, caps["chunks"]),
            raw_rows=list(raw_rows),
        )

    def _normalize_row(self, row: dict) -> dict:
        return {
            "layer": str(row.get("layer") or ""),
            "path": str(row.get("path") or ""),
            "title": str(row.get("title") or ""),
            "content": str(row.get("content") or ""),
            "metadata": dict(row.get("metadata") or {}),
        }

    def _support_paths(self, sub_intent: str, buckets: dict[str, list[dict]]) -> list[str]:
        ordered: list[dict] = []
        if sub_intent == "SYSTEM_FLOW_EXPLAIN":
            ordered = [*buckets["workflows"], *buckets["relations"], *buckets["documents"]]
        elif sub_intent == "COMPONENT_EXPLAIN":
            ordered = [*buckets["facts"], *buckets["relations"], *buckets["documents"]]
        elif sub_intent == "API_METHOD_EXPLAIN":
            ordered = [*buckets["facts"], *buckets["workflows"], *buckets["documents"]]
        elif sub_intent == "ENTITY_EXPLAIN":
            ordered = [*buckets["entities"], *buckets["relations"], *buckets["documents"]]
        elif sub_intent == "RELATED_DOCS_EXPLAIN":
            ordered = [*buckets["relations"], *buckets["documents"], *buckets["entities"]]
        elif sub_intent == "GENERIC_QA":
            ordered = [*buckets["documents"], *buckets["chunks"]]
        else:
            ordered = [*buckets["facts"], *buckets["documents"]]
        paths: list[str] = []
        for item in ordered:
            path = str(item.get("path") or "").strip()
            if path and path not in paths:
                paths.append(path)
        return paths[:4]

    def _select_chunks(self, chunks: list[dict], support_paths: list[str], limit: int) -> list[dict]:
        if limit <= 0:
            return []
        targeted = [item for item in chunks if str(item.get("path") or "") in support_paths]
        selected = targeted[:limit]
        if len(selected) >= limit:
            return selected
        for item in chunks:
            if item in selected:
                continue
            selected.append(item)
            if len(selected) >= limit:
                break
        return selected
