"""Vector store backends for RAG system."""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any

from logos_adk.rag import Chunk


from logos_adk.math_utils import cosine_similarity as _cosine_similarity  # noqa: F401 — re-export for backwards compat


class VectorStore(ABC):
    """Abstract base class for vector stores."""

    @abstractmethod
    async def add(self, chunks: list[Chunk]) -> None:
        ...

    @abstractmethod
    async def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        ...

    @abstractmethod
    async def delete(self, document_id: str) -> None:
        ...

    @abstractmethod
    async def clear(self) -> None:
        ...


class InMemoryVectorStore(VectorStore):
    """In-memory vector store using cosine similarity."""

    def __init__(self) -> None:
        self._chunks: dict[str, Chunk] = {}

    async def add(self, chunks: list[Chunk]) -> None:
        for chunk in chunks:
            self._chunks[chunk.id] = chunk

    async def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        scored: list[tuple[float, Chunk]] = []
        for chunk in self._chunks.values():
            if chunk.embedding is None:
                continue
            if filter:
                if not all(chunk.metadata.get(k) == v for k, v in filter.items()):
                    continue
            sim = _cosine_similarity(query_vector, chunk.embedding)
            scored.append((sim, chunk))
        scored.sort(key=lambda x: x[0], reverse=True)
        return [chunk for _, chunk in scored[:limit]]

    async def delete(self, document_id: str) -> None:
        to_remove = [cid for cid, c in self._chunks.items() if c.document_id == document_id]
        for cid in to_remove:
            del self._chunks[cid]

    async def clear(self) -> None:
        self._chunks.clear()
