"""Vector store implementations for RAG pipelines."""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass, field


@dataclass
class Document:
    """A document with optional embedding."""

    id: str
    text: str
    metadata: dict = field(default_factory=dict)
    embedding: list[float] | None = None


class VectorStore(ABC):
    """Base class for vector stores."""

    @abstractmethod
    def add(self, documents: list[Document]) -> None:
        ...

    @abstractmethod
    def query(self, query_embedding: list[float], top_k: int = 5) -> list[Document]:
        ...

    @abstractmethod
    def delete(self, ids: list[str]) -> None:
        ...

    @abstractmethod
    def count(self) -> int:
        ...


class ChromaStore(VectorStore):
    """ChromaDB vector store implementation.

    Uses an ephemeral in-memory client by default. Pass a persistent path
    for disk-based storage.
    """

    def __init__(
        self,
        collection_name: str = "default",
        persist_directory: str | None = None,
    ):
        import chromadb

        self.collection_name = collection_name
        if persist_directory:
            self._client = chromadb.PersistentClient(path=persist_directory)
        else:
            self._client = chromadb.EphemeralClient()

        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    def add(self, documents: list[Document]) -> None:
        if not documents:
            return

        ids = [doc.id for doc in documents]
        texts = [doc.text for doc in documents]
        metadatas = [doc.metadata if doc.metadata else None for doc in documents]
        embeddings = [doc.embedding for doc in documents if doc.embedding is not None]

        kwargs = {"ids": ids, "documents": texts}
        # ChromaDB rejects empty metadata dicts — only pass if non-empty
        if any(m is not None for m in metadatas):
            kwargs["metadatas"] = [m or {"_": "1"} for m in metadatas]
        if embeddings and len(embeddings) == len(documents):
            kwargs["embeddings"] = embeddings

        self._collection.add(**kwargs)

    def query(self, query_embedding: list[float], top_k: int = 5) -> list[Document]:
        results = self._collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
        )

        documents = []
        if results["ids"] and results["ids"][0]:
            for i, doc_id in enumerate(results["ids"][0]):
                text = results["documents"][0][i] if results["documents"] else ""
                metadata = results["metadatas"][0][i] if results["metadatas"] else {}
                documents.append(Document(id=doc_id, text=text, metadata=metadata))

        return documents

    def delete(self, ids: list[str]) -> None:
        self._collection.delete(ids=ids)

    def count(self) -> int:
        return self._collection.count()
