from __future__ import annotations

from app.modules.rag.contracts import RagLayer, RetrievalMode


class RagQueryRouter:
    _CODE_HINTS = (
        "как работает код",
        "explain code",
        "explain the code",
        "по коду",
        "из кода",
        "построй документацию по коду",
        "документацию по коду",
        "where is implemented",
        "где реализовано",
        "endpoint",
        "handler",
        "symbol",
        "function",
        "class",
        "method",
    )

    _DOCS_LAYERS = [
        RagLayer.DOCS_MODULE_CATALOG,
        RagLayer.DOCS_FACT_INDEX,
        RagLayer.DOCS_SECTION_INDEX,
        RagLayer.DOCS_POLICY_INDEX,
    ]
    _CODE_LAYERS = [
        RagLayer.CODE_ENTRYPOINTS,
        RagLayer.CODE_SYMBOL_CATALOG,
        RagLayer.CODE_DEPENDENCY_GRAPH,
        RagLayer.CODE_SOURCE_CHUNKS,
    ]

    def resolve_mode(self, query: str) -> str:
        lowered = query.lower()
        return RetrievalMode.CODE if any(hint in lowered for hint in self._CODE_HINTS) else RetrievalMode.DOCS

    def layers_for_mode(self, mode: str) -> list[str]:
        return list(self._CODE_LAYERS if mode == RetrievalMode.CODE else self._DOCS_LAYERS)