"""
ChromaDB vector store wrapper.

ChromaDB stores: the raw text, the embedding vector, and metadata (source, page).
On disk so the vector store survives between runs — you don't re-embed on every restart.

Interview concept to own:
  Chroma uses HNSW (Hierarchical Navigable Small World) indexing under the hood.
  HNSW is an approximate nearest neighbour algorithm — it trades a tiny bit of
  recall for massive speed gains at scale. At 10M vectors, brute-force cosine
  search is unusable; HNSW is O(log n).
"""

import chromadb
from chromadb.config import Settings
from src.utils.config import config
from src.utils.logger import logger


class VectorStore:

    def __init__(self, collection_name: str = None, persist_dir: str = None):
        collection_name = collection_name or config.collection_name
        persist_dir = persist_dir or config.chroma_persist_dir

        self.client = chromadb.PersistentClient(
            path=persist_dir,
            settings=Settings(anonymized_telemetry=False),
        )
        self.collection = self.client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # Use cosine distance, not L2
        )
        logger.info(f"Vector store ready: '{collection_name}' at {persist_dir}")
        logger.info(f"  Current document count: {self.collection.count()}")

    def add_chunks(self, chunks: list[dict], embeddings: list[list[float]]) -> None:
        """
        Store chunks + their embeddings in Chroma.

        Args:
            chunks:     Output from chunker.chunk_pages()
            embeddings: Parallel list of embedding vectors from Embedder.embed_texts()
        """
        if not chunks:
            logger.warning("No chunks to add.")
            return

        self.collection.add(
            ids=[c["chunk_id"] for c in chunks],
            embeddings=embeddings,
            documents=[c["text"] for c in chunks],
            metadatas=[{"source": c["source"], "page": c["page"]} for c in chunks],
        )
        logger.info(f"Added {len(chunks)} chunks. Total in store: {self.collection.count()}")

    def query(self, query_embedding: list[float], top_k: int = None) -> list[dict]:
        """
        Find the top-k most similar chunks to a query embedding.

        Returns:
            List of dicts: [{"text": str, "source": str, "page": int, "score": float}]
            Ordered by similarity (most similar first).
        """
        top_k = top_k or config.top_k
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"],
        )

        chunks = []
        for doc, meta, dist in zip(
            results["documents"][0],
            results["metadatas"][0],
            results["distances"][0],
        ):
            chunks.append({
                "text": doc,
                "source": meta["source"],
                "page": meta["page"],
                "score": round(1 - dist, 4),  # Convert cosine distance → similarity
            })

        return chunks

    def reset(self) -> None:
        """Delete and recreate the collection. Use during development to re-ingest."""
        name = self.collection.name
        self.client.delete_collection(name)
        self.collection = self.client.get_or_create_collection(
            name=name,
            metadata={"hnsw:space": "cosine"},
        )
        logger.warning(f"Collection '{name}' has been reset.")