from __future__ import annotations

from app.modules.rag.contracts.enums import RagLayer
from app.modules.rag.intent_router_v2.retrieval.layer_query_builder import LayerQueryBuilder
from app.modules.rag.intent_router_v2.models import ConversationState, QueryAnchor, RepoContext, RetrievalSpec
from app.modules.rag.intent_router_v2.retrieval.retrieval_filter_builder import RetrievalFilterBuilder


class RetrievalSpecFactory:
    _LAYERS = {
        "CODE_QA": [
            (RagLayer.CODE_SYMBOL_CATALOG, 8),
            (RagLayer.CODE_SOURCE_CHUNKS, 8),
            (RagLayer.CODE_SEMANTIC_ROLES, 8),
            (RagLayer.CODE_DEPENDENCY_GRAPH, 6),
            (RagLayer.CODE_ENTRYPOINTS, 6),
        ],
        "DOCS_QA": [
            (RagLayer.DOCS_MODULE_CATALOG, 5),
            (RagLayer.DOCS_FACT_INDEX, 8),
            (RagLayer.DOCS_SECTION_INDEX, 8),
            (RagLayer.DOCS_POLICY_INDEX, 4),
        ],
        "GENERATE_DOCS_FROM_CODE": [
            (RagLayer.CODE_SYMBOL_CATALOG, 12),
            (RagLayer.CODE_DEPENDENCY_GRAPH, 8),
            (RagLayer.CODE_SOURCE_CHUNKS, 12),
            (RagLayer.CODE_ENTRYPOINTS, 6),
        ],
        "PROJECT_MISC": [
            (RagLayer.DOCS_MODULE_CATALOG, 4),
            (RagLayer.DOCS_SECTION_INDEX, 6),
            (RagLayer.CODE_SYMBOL_CATALOG, 4),
            (RagLayer.CODE_SOURCE_CHUNKS, 4),
        ],
    }
    _DOMAINS = {
        "CODE_QA": ["CODE"],
        "DOCS_QA": ["DOCS"],
        "GENERATE_DOCS_FROM_CODE": ["CODE"],
        "PROJECT_MISC": ["CODE", "DOCS"],
    }
    _RERANK = {
        "CODE_QA": "code",
        "DOCS_QA": "docs",
        "GENERATE_DOCS_FROM_CODE": "generate",
        "PROJECT_MISC": "project",
    }
    _OPEN_FILE_LAYERS = [
        (RagLayer.CODE_SOURCE_CHUNKS, 12),
    ]
    _FIND_ENTRYPOINTS_LAYERS = [
        (RagLayer.CODE_ENTRYPOINTS, 12),
        (RagLayer.CODE_SOURCE_CHUNKS, 6),
    ]
    _EXPLAIN_LOCAL_LAYERS = [
        (RagLayer.CODE_SYMBOL_CATALOG, 8),
        (RagLayer.CODE_SOURCE_CHUNKS, 12),
        (RagLayer.CODE_DEPENDENCY_GRAPH, 4),
    ]
    _FIND_TESTS_LAYERS = [
        (RagLayer.CODE_SYMBOL_CATALOG, 8),
        (RagLayer.CODE_DEPENDENCY_GRAPH, 6),
        (RagLayer.CODE_SOURCE_CHUNKS, 10),
    ]
    _TRACE_FLOW_LAYERS = [
        (RagLayer.CODE_SYMBOL_CATALOG, 8),
        (RagLayer.CODE_SOURCE_CHUNKS, 8),
        (RagLayer.CODE_SEMANTIC_ROLES, 8),
        (RagLayer.CODE_DEPENDENCY_GRAPH, 10),
        (RagLayer.CODE_ENTRYPOINTS, 10),
    ]
    _ARCHITECTURE_LAYERS = [
        (RagLayer.CODE_SEMANTIC_ROLES, 12),
        (RagLayer.CODE_ENTRYPOINTS, 8),
        (RagLayer.CODE_DEPENDENCY_GRAPH, 8),
        (RagLayer.CODE_SYMBOL_CATALOG, 6),
        (RagLayer.CODE_SOURCE_CHUNKS, 4),
    ]
    _DOCS_SCOPED_LAYERS = [
        (RagLayer.DOCS_SECTION_INDEX, 8),
        (RagLayer.DOCS_FACT_INDEX, 8),
    ]

    def __init__(
        self,
        layer_builder: LayerQueryBuilder | None = None,
        filter_builder: RetrievalFilterBuilder | None = None,
    ) -> None:
        self._layer_builder = layer_builder or LayerQueryBuilder()
        self._filter_builder = filter_builder or RetrievalFilterBuilder()

    def build(
        self,
        intent: str,
        anchors: list[QueryAnchor],
        repo_context: RepoContext,
        *,
        raw_query: str = "",
        conversation_state: ConversationState | None = None,
        conversation_mode: str = "START",
        sub_intent: str = "EXPLAIN",
    ) -> RetrievalSpec:
        domains = self._domains(intent, repo_context)
        layers_map = self._with_sub_intent_layers(intent, sub_intent, anchors)
        layer_queries = self._layer_builder.build(intent, repo_context, domains=domains, layers_map=layers_map)
        filters = self._filter_builder.build(
            domains,
            anchors,
            repo_context,
            raw_query=raw_query,
            conversation_state=conversation_state,
            conversation_mode=conversation_mode,
            sub_intent=sub_intent,
        )
        if intent == "DOCS_QA" and list(getattr(filters, "path_scope", []) or []):
            scoped_map = dict(self._LAYERS)
            scoped_map["DOCS_QA"] = list(self._DOCS_SCOPED_LAYERS)
            layer_queries = self._layer_builder.build(intent, repo_context, domains=domains, layers_map=scoped_map)
        return RetrievalSpec(
            domains=domains,
            layer_queries=layer_queries,
            filters=filters,
            rerank_profile=self._RERANK[intent],
        )

    def _domains(self, intent: str, repo_context: RepoContext) -> list[str]:
        available = set(repo_context.available_domains or ["CODE", "DOCS"])
        result = [domain for domain in self._DOMAINS[intent] if domain in available]
        return result or list(self._DOMAINS[intent])

    def _with_sub_intent_layers(
        self,
        intent: str,
        sub_intent: str,
        anchors: list[QueryAnchor],
    ) -> dict[str, list[tuple[str, int]]]:
        if intent != "CODE_QA":
            return self._LAYERS
        layers_map = dict(self._LAYERS)
        if sub_intent == "OPEN_FILE":
            layers_map["CODE_QA"] = list(self._OPEN_FILE_LAYERS)
        elif sub_intent == "FIND_ENTRYPOINTS":
            layers_map["CODE_QA"] = list(self._FIND_ENTRYPOINTS_LAYERS)
        elif sub_intent == "FIND_TESTS":
            layers_map["CODE_QA"] = list(self._FIND_TESTS_LAYERS)
        elif sub_intent == "TRACE_FLOW":
            layers_map["CODE_QA"] = list(self._TRACE_FLOW_LAYERS)
        elif sub_intent == "ARCHITECTURE":
            layers_map["CODE_QA"] = list(self._ARCHITECTURE_LAYERS)
        elif sub_intent == "EXPLAIN_LOCAL":
            layers_map["CODE_QA"] = list(self._EXPLAIN_LOCAL_LAYERS)
        elif sub_intent != "EXPLAIN" and not self._needs_entrypoints(anchors):
            layers_map["CODE_QA"] = [
                (RagLayer.CODE_SYMBOL_CATALOG, 8),
                (RagLayer.CODE_SOURCE_CHUNKS, 8),
                (RagLayer.CODE_DEPENDENCY_GRAPH, 6),
            ]
        return layers_map

    def _needs_entrypoints(self, anchors: list[QueryAnchor]) -> bool:
        values = " ".join(anchor.value.lower() for anchor in anchors if anchor.type in {"KEY_TERM", "SYMBOL"})
        markers = ("entrypoint", "endpoint", "вызыва", "поток", "flow", "запуска")
        return any(marker in values for marker in markers)
