from __future__ import annotations

import json

from sqlalchemy import bindparam, text

from app.core.rag.persistence.retrieval_statement_builder import RetrievalStatementBuilder
from app.core.shared.database import get_engine


class RagQueryRepository:
    def __init__(self) -> None:
        self._builder = RetrievalStatementBuilder()

    def retrieve(
        self,
        rag_session_id: str,
        query_embedding: list[float],
        *,
        query_text: str = "",
        limit: int = 5,
        layers: list[str] | None = None,
        path_prefixes: list[str] | None = None,
        exclude_path_prefixes: list[str] | None = None,
        exclude_like_patterns: list[str] | None = None,
        prefer_path_prefixes: list[str] | None = None,
        prefer_like_patterns: list[str] | None = None,
        metadata_domain: str | None = None,
        metadata_subdomain: str | None = None,
        metadata_type: str | None = None,
        metadata_doc_type: str | None = None,
        prefer_non_tests: bool = False,
    ) -> list[dict]:
        sql, params = self._builder.build_retrieve(
            rag_session_id,
            query_embedding,
            query_text=query_text,
            limit=limit,
            layers=layers,
            path_prefixes=path_prefixes,
            exclude_path_prefixes=exclude_path_prefixes,
            exclude_like_patterns=exclude_like_patterns,
            prefer_path_prefixes=prefer_path_prefixes,
            prefer_like_patterns=prefer_like_patterns,
            metadata_domain=metadata_domain,
            metadata_subdomain=metadata_subdomain,
            metadata_type=metadata_type,
            metadata_doc_type=metadata_doc_type,
            prefer_non_tests=prefer_non_tests,
        )
        with get_engine().connect() as conn:
            rows = conn.execute(text(sql), params).mappings().fetchall()
        return [self._row_to_dict(row) for row in rows]

    def retrieve_lexical_code(
        self,
        rag_session_id: str,
        *,
        query_text: str,
        limit: int = 5,
        path_prefixes: list[str] | None = None,
        exclude_path_prefixes: list[str] | None = None,
        exclude_like_patterns: list[str] | None = None,
        prefer_path_prefixes: list[str] | None = None,
        prefer_like_patterns: list[str] | None = None,
        prefer_non_tests: bool = False,
    ) -> list[dict]:
        sql, params = self._builder.build_lexical_code(
            rag_session_id,
            query_text=query_text,
            limit=limit,
            path_prefixes=path_prefixes,
            exclude_path_prefixes=exclude_path_prefixes,
            exclude_like_patterns=exclude_like_patterns,
            prefer_path_prefixes=prefer_path_prefixes,
            prefer_like_patterns=prefer_like_patterns,
            prefer_non_tests=prefer_non_tests,
        )
        if sql is None:
            return []
        with get_engine().connect() as conn:
            rows = conn.execute(text(sql), params).mappings().fetchall()
        return [self._row_to_dict(row) for row in 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,
    ) -> list[dict]:
        normalized_paths = [str(path).strip() for path in paths if str(path).strip()]
        if not normalized_paths:
            return []
        filters = ["rag_session_id = :sid", "path IN :paths"]
        params: dict = {
            "sid": rag_session_id,
            "paths": normalized_paths,
            "lim": max(1, int(limit)),
        }
        if repo_id:
            filters.append("repo_id = :repo_id")
            params["repo_id"] = repo_id
        stmt = text(
            f"""
            SELECT path, content, layer, title, metadata_json, span_start, span_end,
                   0 AS lexical_rank,
                   0 AS prefer_bonus,
                   0 AS test_penalty,
                   0 AS structural_rank,
                   0 AS layer_rank,
                   0 AS distance
            FROM rag_chunks
            WHERE {' AND '.join(filters)}
            ORDER BY path ASC, COALESCE(span_start, 0) ASC, COALESCE(chunk_index, 0) ASC
            LIMIT :lim
            """
        ).bindparams(bindparam("paths", expanding=True))
        if layers:
            params["layers"] = [str(item).strip() for item in layers if str(item).strip()]
            if params["layers"]:
                stmt = text(
                    f"""
                    SELECT path, content, layer, title, metadata_json, span_start, span_end,
                           0 AS lexical_rank,
                           0 AS prefer_bonus,
                           0 AS test_penalty,
                           0 AS structural_rank,
                           0 AS layer_rank,
                           0 AS distance
                    FROM rag_chunks
                    WHERE {' AND '.join(filters)} AND layer IN :layers
                    ORDER BY path ASC, COALESCE(span_start, 0) ASC, COALESCE(chunk_index, 0) ASC
                    LIMIT :lim
                    """
                ).bindparams(
                    bindparam("paths", expanding=True),
                    bindparam("layers", expanding=True),
                )
        with get_engine().connect() as conn:
            rows = conn.execute(stmt, params).mappings().fetchall()
        return [self._row_to_dict(row) for row in rows]

    def retrieve_by_document_ids(
        self,
        rag_session_id: str,
        *,
        document_ids: list[str],
        layers: list[str] | None = None,
        limit: int = 200,
    ) -> list[dict]:
        normalized_ids = [str(item).strip() for item in document_ids if str(item).strip()]
        if not normalized_ids:
            return []
        params: dict = {
            "sid": rag_session_id,
            "lim": max(1, int(limit)),
            "doc_ids_document": normalized_ids,
            "doc_ids_doc": normalized_ids,
            "doc_ids_subject": normalized_ids,
        }
        stmt = text(
            """
            SELECT path, content, layer, title, metadata_json, span_start, span_end,
                   0 AS lexical_rank,
                   0 AS prefer_bonus,
                   0 AS test_penalty,
                   0 AS structural_rank,
                   0 AS layer_rank,
                   0 AS distance
            FROM rag_chunks
            WHERE rag_session_id = :sid
              AND (
                COALESCE(metadata_json->>'document_id', '') IN :doc_ids_document
                OR COALESCE(metadata_json->>'doc_id', '') IN :doc_ids_doc
                OR COALESCE(metadata_json->>'subject_id', '') IN :doc_ids_subject
              )
            ORDER BY
              CASE layer
                WHEN 'D1_DOCUMENT_CATALOG' THEN 0
                WHEN 'D2_FACT_INDEX' THEN 1
                WHEN 'D0_DOC_CHUNKS' THEN 2
                WHEN 'D4_WORKFLOW_INDEX' THEN 3
                WHEN 'D6_INTEGRATION_INDEX' THEN 4
                WHEN 'D5_RELATION_GRAPH' THEN 5
                ELSE 9
              END,
              path ASC,
              COALESCE(span_start, 0) ASC,
              COALESCE(chunk_index, 0) ASC
            LIMIT :lim
            """
        ).bindparams(
            bindparam("doc_ids_document", expanding=True),
            bindparam("doc_ids_doc", expanding=True),
            bindparam("doc_ids_subject", expanding=True),
        )
        if layers:
            params["layers"] = [str(item).strip() for item in layers if str(item).strip()]
            if params["layers"]:
                stmt = text(
                    """
                    SELECT path, content, layer, title, metadata_json, span_start, span_end,
                           0 AS lexical_rank,
                           0 AS prefer_bonus,
                           0 AS test_penalty,
                           0 AS structural_rank,
                           0 AS layer_rank,
                           0 AS distance
                    FROM rag_chunks
                    WHERE rag_session_id = :sid
                      AND layer IN :layers
                      AND (
                        COALESCE(metadata_json->>'document_id', '') IN :doc_ids_document
                        OR COALESCE(metadata_json->>'doc_id', '') IN :doc_ids_doc
                        OR COALESCE(metadata_json->>'subject_id', '') IN :doc_ids_subject
                      )
                    ORDER BY
                      CASE layer
                        WHEN 'D1_DOCUMENT_CATALOG' THEN 0
                        WHEN 'D2_FACT_INDEX' THEN 1
                        WHEN 'D0_DOC_CHUNKS' THEN 2
                        WHEN 'D4_WORKFLOW_INDEX' THEN 3
                        WHEN 'D6_INTEGRATION_INDEX' THEN 4
                        WHEN 'D5_RELATION_GRAPH' THEN 5
                        ELSE 9
                      END,
                      path ASC,
                      COALESCE(span_start, 0) ASC,
                      COALESCE(chunk_index, 0) ASC
                    LIMIT :lim
                    """
                ).bindparams(
                    bindparam("layers", expanding=True),
                    bindparam("doc_ids_document", expanding=True),
                    bindparam("doc_ids_doc", expanding=True),
                    bindparam("doc_ids_subject", expanding=True),
                )
        with get_engine().connect() as conn:
            rows = conn.execute(stmt, params).mappings().fetchall()
        return [self._row_to_dict(row) for row in rows]

    def retrieve_chunks_by_path_substrings(
        self,
        rag_session_id: str,
        *,
        path_needles: list[str],
        layers: list[str] | None = None,
        limit: int = 200,
    ) -> list[dict]:
        normalized_needles = [str(item).strip().lower() for item in path_needles if str(item).strip()]
        if not normalized_needles:
            return []
        params: dict = {
            "sid": rag_session_id,
            "lim": max(1, int(limit)),
        }
        filters = ["rag_session_id = :sid"]
        like_parts: list[str] = []
        for idx, needle in enumerate(normalized_needles):
            key = f"needle_{idx}"
            params[key] = f"%{needle}%"
            like_parts.append(f"lower(path) LIKE :{key}")
        filters.append("(" + " OR ".join(like_parts) + ")")
        if layers:
            normalized_layers = [str(item).strip() for item in layers if str(item).strip()]
            if normalized_layers:
                params["layers"] = normalized_layers
                filters.append("layer IN :layers")
        stmt = text(
            f"""
            SELECT path, content, layer, title, metadata_json, span_start, span_end,
                   0 AS lexical_rank,
                   0 AS prefer_bonus,
                   0 AS test_penalty,
                   0 AS structural_rank,
                   0 AS layer_rank,
                   0 AS distance
            FROM rag_chunks
            WHERE {' AND '.join(filters)}
            ORDER BY path ASC, COALESCE(span_start, 0) ASC, COALESCE(chunk_index, 0) ASC
            LIMIT :lim
            """
        )
        if "layers" in params:
            stmt = stmt.bindparams(bindparam("layers", expanding=True))
        with get_engine().connect() as conn:
            rows = conn.execute(stmt, params).mappings().fetchall()
        return [self._row_to_dict(row) for row in rows]

    def list_docs_scope_index_rows(
        self,
        rag_session_id: str,
        *,
        limit: int = 8000,
    ) -> list[dict]:
        """Return D1/D3 catalog rows for a session (no vector retrieval, no D0 chunks).

        Used by the v2 intent router to derive scope candidates from indexed DOCS layers only.
        """
        sid = str(rag_session_id or "").strip()
        if not sid:
            return []
        lim = max(1, min(int(limit), 50_000))
        layers = ("D1_DOCUMENT_CATALOG", "D3_ENTITY_CATALOG")
        stmt = text(
            """
            SELECT path, content, layer, title, metadata_json, span_start, span_end,
                   0 AS lexical_rank,
                   0 AS prefer_bonus,
                   0 AS test_penalty,
                   0 AS structural_rank,
                   0 AS layer_rank,
                   0 AS distance
            FROM rag_chunks
            WHERE rag_session_id = :sid
              AND layer IN :layers
            ORDER BY layer ASC, path ASC, COALESCE(span_start, 0) ASC, COALESCE(chunk_index, 0) ASC
            LIMIT :lim
            """
        ).bindparams(bindparam("layers", expanding=True))
        params: dict = {"sid": sid, "layers": list(layers), "lim": lim}
        with get_engine().connect() as conn:
            rows = conn.execute(stmt, params).mappings().fetchall()
        return [self._row_to_dict(row) for row in rows]

    def _row_to_dict(self, row) -> dict:
        data = dict(row)
        raw_metadata = data.pop("metadata_json")
        if isinstance(raw_metadata, dict):
            data["metadata"] = raw_metadata
            return data
        data["metadata"] = json.loads(str(raw_metadata or "{}"))
        return data
