from __future__ import annotations

import json

from sqlalchemy import bindparam, text

from app.modules.rag.persistence.retrieval_statement_builder import RetrievalStatementBuilder
from app.modules.shared.db 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,
        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,
            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 _row_to_dict(self, row) -> dict:
        data = dict(row)
        data["metadata"] = json.loads(str(data.pop("metadata_json") or "{}"))
        return data
