from __future__ import annotations

import json
from datetime import date, datetime

from sqlalchemy import text

from app.modules.rag.contracts import RagDocument


class RagDocumentRepository:
    def replace_documents(self, conn, rag_session_id: str, docs: list[RagDocument]) -> None:
        conn.execute(text("DELETE FROM rag_chunks WHERE rag_session_id = :sid"), {"sid": rag_session_id})
        self.insert_documents(conn, rag_session_id, docs)

    def apply_document_changes(
        self,
        conn,
        rag_session_id: str,
        delete_paths: list[str],
        docs: list[RagDocument],
    ) -> None:
        self._delete_paths(conn, rag_session_id, delete_paths)
        if not docs:
            return
        self._delete_paths(conn, rag_session_id, sorted({doc.source.path for doc in docs}))
        self.insert_documents(conn, rag_session_id, docs)

    def insert_documents(self, conn, rag_session_id: str, docs: list[RagDocument]) -> None:
        for doc in docs:
            conn.execute(text(self._insert_sql()), self._insert_params(rag_session_id, doc))

    def _delete_paths(self, conn, rag_session_id: str, paths: list[str]) -> None:
        if not paths:
            return
        conn.execute(
            text("DELETE FROM rag_chunks WHERE rag_session_id = :sid AND path = ANY(:paths)"),
            {"sid": rag_session_id, "paths": paths},
        )

    def _insert_sql(self) -> str:
        return """
            INSERT INTO rag_chunks (
                rag_session_id, path, chunk_index, content, embedding, layer, lang, repo_id,
                commit_sha, title, metadata_json, span_start, span_end, created_at, updated_at
            )
            VALUES (
                :sid, :path, :chunk_index, :content, CAST(:emb AS vector), :layer, :lang, :repo_id,
                :commit_sha, :title, CAST(:metadata_json AS jsonb), :span_start, :span_end,
                CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
            )
        """

    def _insert_params(self, rag_session_id: str, doc: RagDocument) -> dict:
        row = doc.to_record()
        metadata = self._json_ready(row["metadata"])
        emb = row["embedding"] or []
        return {
            "sid": rag_session_id,
            "path": row["path"],
            "chunk_index": int(metadata.get("chunk_index", 0)),
            "content": row["text"],
            "emb": "[" + ",".join(str(x) for x in emb) + "]" if emb else None,
            "layer": row["layer"],
            "lang": row["lang"],
            "repo_id": row["repo_id"],
            "commit_sha": row["commit_sha"],
            "title": row["title"],
            "metadata_json": json.dumps(metadata, ensure_ascii=True),
            "span_start": row["span_start"],
            "span_end": row["span_end"],
        }

    def _json_ready(self, value):
        if isinstance(value, dict):
            return {str(key): self._json_ready(item) for key, item in value.items()}
        if isinstance(value, list):
            return [self._json_ready(item) for item in value]
        if isinstance(value, tuple):
            return [self._json_ready(item) for item in value]
        if isinstance(value, (date, datetime)):
            return value.isoformat()
        return value
