"""ChromaDB vector store wrapper — embeds, stores, and searches text chunks."""

from __future__ import annotations

import os
import uuid
from pathlib import Path
from typing import Optional

import chromadb
from chromadb.config import Settings


class VectorStore:
    """Wrapper around ChromaDB for chunk storage and similarity search.

    Uses ChromaDB's built-in embedding (all-MiniLM-L6-v2) by default,
    or a custom embedding function.
    """

    def __init__(
        self,
        data_dir: str | Path = "./res_sum_data",
        collection_name: str = "paper_chunks",
        embedding_function: Optional[object] = None,
    ) -> None:
        """Initialize the vector store.

        Args:
            data_dir: Directory for persistent ChromaDB storage.
            collection_name: Name of the ChromaDB collection.
            embedding_function: Custom ChromaDB-compatible embedding function.
                If None, uses ChromaDB's default (all-MiniLM-L6-v2).
        """
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)

        chroma_path = self.data_dir / "chroma_db"
        self._client = chromadb.PersistentClient(
            path=str(chroma_path),
            settings=Settings(anonymized_telemetry=False),
        )

        kwargs = {"name": collection_name}
        if embedding_function is not None:
            kwargs["embedding_function"] = embedding_function

        self._collection = self._client.get_or_create_collection(**kwargs)

    @property
    def collection(self) -> chromadb.Collection:
        return self._collection

    def add_chunks(
        self,
        chunks: list[str],
        metadatas: Optional[list[dict]] = None,
        ids: Optional[list[str]] = None,
    ) -> list[str]:
        """Add text chunks to the vector store.

        Args:
            chunks: List of text chunks.
            metadatas: Optional metadata for each chunk.
            ids: Optional IDs. If None, UUIDs are generated.

        Returns:
            List of chunk IDs.
        """
        if not chunks:
            return []

        if ids is None:
            ids = [str(uuid.uuid4()) for _ in chunks]

        if metadatas is None:
            metadatas = [{"source": "res_sum"} for _ in chunks]

        # ChromaDB has a batch size limit; add in batches of 5000
        batch_size = 5000
        for i in range(0, len(chunks), batch_size):
            batch_end = min(i + batch_size, len(chunks))
            self._collection.add(
                documents=chunks[i:batch_end],
                metadatas=metadatas[i:batch_end],
                ids=ids[i:batch_end],
            )

        return ids

    def search(
        self,
        query: str,
        n_results: int = 10,
        where: Optional[dict] = None,
    ) -> dict:
        """Search for similar chunks.

        Args:
            query: Query text.
            n_results: Number of results to return.
            where: Optional metadata filter.

        Returns:
            ChromaDB query results dict with keys: ids, documents, metadatas, distances.
        """
        kwargs = {
            "query_texts": [query],
            "n_results": min(n_results, self._collection.count() or 1),
        }
        if where:
            kwargs["where"] = where

        return self._collection.query(**kwargs)

    def get_chunks_by_ids(self, ids: list[str]) -> dict:
        """Retrieve specific chunks by their IDs."""
        return self._collection.get(ids=ids)

    def count(self) -> int:
        """Return the number of chunks in the store."""
        return self._collection.count()

    def has_paper(self, paper_filename: str) -> bool:
        """Check if a paper has already been ingested."""
        results = self._collection.get(
            where={"paper_filename": paper_filename},
            limit=1,
        )
        return len(results["ids"]) > 0

    def delete_paper(self, paper_filename: str) -> None:
        """Delete all chunks for a specific paper."""
        results = self._collection.get(
            where={"paper_filename": paper_filename},
        )
        if results["ids"]:
            self._collection.delete(ids=results["ids"])

    def list_papers(self) -> list[str]:
        """List all ingested paper filenames."""
        all_data = self._collection.get(include=["metadatas"])
        papers = set()
        for meta in all_data.get("metadatas", []):
            if meta and "paper_filename" in meta:
                papers.add(meta["paper_filename"])
        return sorted(papers)
