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

from typing import Any

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


class ChromaVectorStore(VectorStore):
    """Vector store backed by ChromaDB."""

    def __init__(
        self,
        collection_name: str = "logos_rag",
        client: Any = None,
        path: str | None = None,
    ) -> None:
        self._collection_name = collection_name
        self._client = client
        self._path = path
        self._collection: Any = None

    def _ensure_collection(self) -> Any:
        if self._collection is not None:
            return self._collection
        try:
            import chromadb
        except ImportError:
            raise ImportError(
                "chromadb is required. Install with: pip install logos-adk[rag-chroma]"
            )
        if self._client is None:
            if self._path:
                self._client = chromadb.PersistentClient(path=self._path)
            else:
                self._client = chromadb.Client()
        self._collection = self._client.get_or_create_collection(self._collection_name)
        return self._collection

    async def add(self, chunks: list[Chunk]) -> None:
        collection = self._ensure_collection()
        if not chunks:
            return
        ids = [c.id for c in chunks]
        embeddings = [c.embedding for c in chunks if c.embedding is not None]
        documents = [c.content for c in chunks]
        metadatas = [{**c.metadata, "_document_id": c.document_id} for c in chunks]
        kwargs: dict[str, Any] = {"ids": ids, "documents": documents, "metadatas": metadatas}
        if embeddings and len(embeddings) == len(chunks):
            kwargs["embeddings"] = embeddings
        collection.add(**kwargs)

    async def search(
        self,
        query_vector: list[float],
        limit: int = 5,
        filter: dict[str, Any] | None = None,
    ) -> list[Chunk]:
        collection = self._ensure_collection()
        kwargs: dict[str, Any] = {
            "query_embeddings": [query_vector],
            "n_results": limit,
        }
        if filter:
            kwargs["where"] = filter
        results = collection.query(**kwargs)
        chunks: list[Chunk] = []
        if results and results["ids"]:
            for i, cid in enumerate(results["ids"][0]):
                meta = dict(results["metadatas"][0][i]) if results["metadatas"] else {}
                doc_id = meta.pop("_document_id", "")
                embedding = results["embeddings"][0][i] if results.get("embeddings") else None
                chunks.append(Chunk(
                    content=results["documents"][0][i] if results["documents"] else "",
                    metadata=meta,
                    document_id=doc_id,
                    embedding=embedding,
                    id=cid,
                ))
        return chunks

    async def delete(self, document_id: str) -> None:
        collection = self._ensure_collection()
        collection.delete(where={"_document_id": document_id})

    async def clear(self) -> None:
        if self._client is not None:
            try:
                self._client.delete_collection(self._collection_name)
            except Exception:
                pass
        self._collection = None
