from __future__ import annotations

from time import perf_counter
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from app.modules.rag.persistence.repository import RagRepository


class SessionEmbeddingDimensions:
    def __init__(self, repository: RagRepository | None = None) -> None:
        self._cache: dict[str, int] = {}
        self._repository = repository

    def resolve(self, rag_session_id: str) -> int:
        if rag_session_id in self._cache:
            return self._cache[rag_session_id]
        from app.modules.shared.db import get_engine
        from sqlalchemy import text

        with get_engine().connect() as conn:
            row = conn.execute(
                text(
                    """
                    SELECT vector_dims(embedding) AS dim
                    FROM rag_chunks
                    WHERE rag_session_id = :sid AND embedding IS NOT NULL
                    LIMIT 1
                    """
                ),
                {"sid": rag_session_id},
            ).mappings().first()
        dim = int(row["dim"]) if row and row.get("dim") else 0
        self._cache[rag_session_id] = dim
        return dim


class CodeQaRetrievalAdapter:
    def __init__(self, repository: RagRepository | None = None) -> None:
        if repository is None:
            from app.modules.rag.persistence.repository import RagRepository

            repository = RagRepository()
        self._repository = repository
        self._dim_resolver = SessionEmbeddingDimensions(self._repository)
        self._last_report: dict[str, Any] | None = None

    def retrieve_with_plan(
        self,
        rag_session_id: str,
        query: str,
        retrieval_spec,
        retrieval_constraints=None,
        *,
        query_plan=None,
    ) -> list[dict]:
        rows: list[dict] = []
        executed_layers: list[str] = []
        per_layer_ms: dict[str, int] = {}
        for layer_query in list(getattr(retrieval_spec, "layer_queries", []) or []):
            layer_id = str(layer_query.layer_id)
            executed_layers.append(layer_id)
            started = perf_counter()
            layer_rows = self._retrieve_layer(
                rag_session_id=rag_session_id,
                query=query,
                layer_id=layer_id,
                limit=max(1, int(layer_query.top_k)),
                path_scope=list(getattr(retrieval_spec.filters, "path_scope", []) or []),
                retrieval_constraints=retrieval_constraints,
                include_tests=str(getattr(retrieval_spec.filters, "test_policy", "EXCLUDE") or "EXCLUDE") == "INCLUDE",
            )
            per_layer_ms[layer_id] = int((perf_counter() - started) * 1000)
            rows.extend(layer_rows)
        self._last_report = {
            "executed_layers": executed_layers,
            "retrieval_mode_by_layer": {layer_id: "vector" for layer_id in executed_layers},
            "top_k_by_layer": {str(item.layer_id): int(item.top_k) for item in list(getattr(retrieval_spec, "layer_queries", []) or [])},
            "filters_by_layer": {
                layer_id: {"path_scope": list(getattr(retrieval_spec.filters, "path_scope", []) or [])}
                for layer_id in executed_layers
            },
            "fallback": {"used": False, "reason": None},
            "retrieval_by_layer_ms": per_layer_ms,
        }
        return self._dedupe(rows)

    def retrieve_exact_files(
        self,
        rag_session_id: str,
        *,
        repo_id: str | None = None,
        paths: list[str],
        layers: list[str] | None = None,
        limit: int = 200,
        query: str = "",
        ranking_profile: str = "",
    ) -> list[dict]:
        started = perf_counter()
        rows = self._repository.retrieve_exact_files(rag_session_id, repo_id=repo_id, paths=paths, layers=layers, limit=limit)
        layer_id = list(layers or ["C0_SOURCE_CHUNKS"])[0]
        self._last_report = {
            "executed_layers": [layer_id],
            "retrieval_mode_by_layer": {layer_id: "exact_path_fetch"},
            "top_k_by_layer": {layer_id: int(limit)},
            "filters_by_layer": {layer_id: {"path_scope": list(paths)}},
            "requests": [{"layer": layer_id, "query": query, "ranking_profile": ranking_profile}],
            "fallback": {"used": False, "reason": None},
            "retrieval_by_layer_ms": {layer_id: int((perf_counter() - started) * 1000)},
        }
        return self._dedupe(rows)

    def hydrate_resolved_symbol_sources(
        self,
        rag_session_id: str,
        base_query: str,
        rag_rows: list[dict],
        symbol_resolution: dict,
        retrieval_spec,
        retrieval_constraints=None,
    ) -> list[dict]:
        if str(symbol_resolution.get("status") or "") != "resolved":
            return rag_rows
        target = str(symbol_resolution.get("resolved_symbol") or "").strip()
        if not target:
            return rag_rows
        symbol_row = next(
            (
                row
                for row in rag_rows
                if str(row.get("layer") or "") == "C1_SYMBOL_CATALOG"
                and target in {str(row.get("title") or "").strip(), str(dict(row.get("metadata") or {}).get("qname") or "").strip()}
            ),
            None,
        )
        if symbol_row is None:
            return rag_rows
        extra = self.retrieve_exact_files(
            rag_session_id,
            paths=[str(symbol_row.get("path") or "")],
            layers=["C0_SOURCE_CHUNKS"],
            limit=6,
            query=base_query,
            ranking_profile=str(getattr(retrieval_spec, "rerank_profile", "") or ""),
        )
        return self._dedupe([*rag_rows, *extra])

    def force_symbol_context_c0(
        self,
        rag_session_id: str,
        *,
        rag_rows: list[dict],
        symbol_resolution: dict,
        limit: int = 20,
    ) -> list[dict]:
        if str(symbol_resolution.get("status") or "") != "resolved":
            return rag_rows
        target = str(symbol_resolution.get("resolved_symbol") or "").strip()
        target_row = next((row for row in rag_rows if str(row.get("title") or "").strip() == target), None)
        if target_row is None:
            return rag_rows
        c0_rows = self.retrieve_exact_files(
            rag_session_id,
            paths=[str(target_row.get("path") or "")],
            layers=["C0_SOURCE_CHUNKS"],
            limit=limit,
            query=target,
            ranking_profile="symbol_context",
        )
        preserved = [row for row in rag_rows if str(row.get("layer") or "") != "C0_SOURCE_CHUNKS"]
        return self._dedupe([*preserved, *c0_rows])

    def consume_retrieval_report(self) -> dict[str, Any] | None:
        report = self._last_report
        self._last_report = None
        return report

    def _retrieve_layer(
        self,
        *,
        rag_session_id: str,
        query: str,
        layer_id: str,
        limit: int,
        path_scope: list[str],
        retrieval_constraints,
        include_tests: bool,
    ) -> list[dict]:
        exclude_prefixes, exclude_like = self._test_filters(include_tests)
        if retrieval_constraints is not None:
            exclude_globs = [item.lower() for item in list(getattr(retrieval_constraints, "exclude_globs", []) or [])]
            exclude_like.extend(self._glob_to_like(item) for item in exclude_globs if "*" in item)
        dim = self._dim_resolver.resolve(rag_session_id)
        if dim <= 0:
            if layer_id != "C0_SOURCE_CHUNKS":
                return []
            return self._repository.retrieve_lexical_code(
                rag_session_id,
                query_text=query,
                limit=limit,
                path_prefixes=path_scope or None,
                exclude_path_prefixes=exclude_prefixes or None,
                exclude_like_patterns=exclude_like or None,
                prefer_non_tests=not include_tests,
            )
        return self._repository.retrieve(
            rag_session_id,
            [0.0] * dim,
            query_text=query,
            limit=limit,
            layers=[layer_id],
            path_prefixes=path_scope or None,
            exclude_path_prefixes=exclude_prefixes or None,
            exclude_like_patterns=exclude_like or None,
            prefer_non_tests=not include_tests,
        )

    def _test_filters(self, include_tests: bool) -> tuple[list[str], list[str]]:
        if include_tests:
            return [], []
        from app.modules.rag.retrieval.test_filter import build_test_filters

        filters = build_test_filters()
        return list(filters.exclude_path_prefixes), list(filters.exclude_like_patterns)

    def _glob_to_like(self, glob: str) -> str:
        return glob.replace("**/", "%/").replace("**", "%").replace("*", "%").replace("?", "_")

    def _dedupe(self, rows: list[dict]) -> list[dict]:
        result: list[dict] = []
        seen: set[tuple[str, str, str, int | None, int | None]] = set()
        for row in rows:
            path = str(row.get("path") or "")
            layer = str(row.get("layer") or "")
            title = str(row.get("title") or "")
            key = (layer, path, title, row.get("span_start"), row.get("span_end"))
            if key in seen:
                continue
            seen.add(key)
            result.append(row)
        return result
