"""PostgreSQL-backed vector store for RAG system."""
from __future__ import annotations

import asyncio
import json
from typing import Any

from logos_adk.postgres_utils import pooled_connect, resolve_database_url, safe_identifier
from logos_adk.rag import Chunk
from logos_adk.rag.vector_store import VectorStore, _cosine_similarity


class PostgreSQLVectorStore(VectorStore):
    """Persistent vector store backed by PostgreSQL."""

    def __init__(self, dsn: str | None = None, table_name: str = "rag_chunks") -> None:
        self.dsn = resolve_database_url(dsn)
        if not self.dsn:
            raise ValueError(
                "PostgreSQL vector store requires a DSN or DATABASE_URL/LOGOS_ADK_DATABASE_URL"
            )
        self.table_name = safe_identifier(table_name)

    def _connect(self) -> Any:
        return pooled_connect(self.dsn)

    def _ensure_schema(self, connection: Any) -> None:
        with connection.cursor() as cursor:
            cursor.execute(
                f"""
                CREATE TABLE IF NOT EXISTS {self.table_name} (
                    id TEXT PRIMARY KEY,
                    document_id TEXT NOT NULL,
                    content TEXT NOT NULL,
                    metadata JSONB NOT NULL,
                    embedding JSONB
                )
                """
            )
            cursor.execute(
                f"CREATE INDEX IF NOT EXISTS idx_{self.table_name}_doc "
                f"ON {self.table_name} (document_id)"
            )
        connection.commit()

    def _add_sync(self, chunks: list[Chunk]) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                for chunk in chunks:
                    cursor.execute(
                        f"""
                        INSERT INTO {self.table_name}
                        (id, document_id, content, metadata, embedding)
                        VALUES (%s, %s, %s, %s::jsonb, %s::jsonb)
                        ON CONFLICT (id) DO UPDATE SET
                            document_id = EXCLUDED.document_id,
                            content = EXCLUDED.content,
                            metadata = EXCLUDED.metadata,
                            embedding = EXCLUDED.embedding
                        """,
                        (
                            chunk.id,
                            chunk.document_id,
                            chunk.content,
                            json.dumps(chunk.metadata),
                            json.dumps(chunk.embedding) if chunk.embedding else None,
                        ),
                    )
            connection.commit()

    async def add(self, chunks: list[Chunk]) -> None:
        await asyncio.to_thread(self._add_sync, chunks)

    def _search_sync(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"""
                    SELECT id, document_id, content, metadata, embedding
                    FROM {self.table_name}
                    WHERE embedding IS NOT NULL
                    """
                )
                rows = cursor.fetchall()

        scored: list[tuple[float, Chunk]] = []
        for row in rows:
            cid, doc_id, content, meta_json, emb_json = row
            meta = meta_json if isinstance(meta_json, dict) else json.loads(meta_json)
            embedding = emb_json if isinstance(emb_json, list) else json.loads(emb_json)
            if filter and not all(meta.get(key) == value for key, value in filter.items()):
                continue
            sim = _cosine_similarity(query_vector, embedding)
            scored.append(
                (
                    sim,
                    Chunk(
                        content=content,
                        metadata=meta,
                        document_id=doc_id,
                        embedding=embedding,
                        id=cid,
                    ),
                )
            )
        scored.sort(key=lambda item: item[0], reverse=True)
        return [chunk for _, chunk in scored[:limit]]

    async def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        return await asyncio.to_thread(self._search_sync, query_vector, limit, filter)

    def _delete_sync(self, document_id: str) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(
                    f"DELETE FROM {self.table_name} WHERE document_id = %s",
                    (document_id,),
                )
            connection.commit()

    async def delete(self, document_id: str) -> None:
        await asyncio.to_thread(self._delete_sync, document_id)

    def _clear_sync(self) -> None:
        with self._connect() as connection:
            self._ensure_schema(connection)
            with connection.cursor() as cursor:
                cursor.execute(f"DELETE FROM {self.table_name}")
            connection.commit()

    async def clear(self) -> None:
        await asyncio.to_thread(self._clear_sync)
