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

import json
import sqlite3
from typing import Any

from logos_adk.rag import Chunk
from logos_adk.rag.vector_store import VectorStore, _cosine_similarity


class SQLiteVectorStore(VectorStore):
    """Persistent vector store backed by SQLite."""

    def __init__(self, path: str) -> None:
        self._path = path
        self._conn = sqlite3.connect(path)
        self._conn.execute("""
            CREATE TABLE IF NOT EXISTS chunks (
                id TEXT PRIMARY KEY,
                document_id TEXT NOT NULL,
                content TEXT NOT NULL,
                metadata TEXT NOT NULL,
                embedding TEXT
            )
        """)
        self._conn.execute(
            "CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(document_id)"
        )
        self._conn.commit()

    async def add(self, chunks: list[Chunk]) -> None:
        for chunk in chunks:
            self._conn.execute(
                "INSERT OR REPLACE INTO chunks (id, document_id, content, metadata, embedding) "
                "VALUES (?, ?, ?, ?, ?)",
                (
                    chunk.id,
                    chunk.document_id,
                    chunk.content,
                    json.dumps(chunk.metadata),
                    json.dumps(chunk.embedding) if chunk.embedding else None,
                ),
            )
        self._conn.commit()

    async def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        cursor = self._conn.execute(
            "SELECT id, document_id, content, metadata, embedding FROM chunks "
            "WHERE embedding IS NOT NULL"
        )
        scored: list[tuple[float, Chunk]] = []
        for row in cursor:
            cid, doc_id, content, meta_json, emb_json = row
            meta = json.loads(meta_json)
            embedding = json.loads(emb_json)
            if filter:
                if not all(meta.get(k) == v for k, v 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 x: x[0], reverse=True)
        return [chunk for _, chunk in scored[:limit]]

    async def delete(self, document_id: str) -> None:
        self._conn.execute("DELETE FROM chunks WHERE document_id = ?", (document_id,))
        self._conn.commit()

    async def clear(self) -> None:
        self._conn.execute("DELETE FROM chunks")
        self._conn.commit()
