"""
Vector store abstraction using ChromaDB.

Provides methods to add documents with embeddings and perform similarity search.
"""
import os
from pathlib import Path
from typing import List, Dict, Any

from chromadb import Client as ChromaClient
from chromadb.config import Settings
from langchain_ollama import OllamaEmbeddings

# Initialize global client (in-memory for simplicity)
client = ChromaClient(Settings(chroma_db_impl="duckdb+parquet", persist_directory=None))
collection_name = "rag_collection"

# Ensure collection exists
if collection_name not in client.list_collections():
    client.create_collection(name=collection_name)
col = client.get_or_create_collection(name=collection_name)

embeddings = OllamaEmbeddings(model="nomic-embed-text")

class ChromaVectorStore:
    """Wrapper around a Chroma collection."""

    def __init__(self, collection):
        self.collection = collection

    def add_documents(self, documents: List[Dict[str, Any]]):
        ids = []
        texts = []
        metadatas = []
        for doc in documents:
            ids.append(doc.get("id", os.urandom(8).hex()))
            texts.append(doc["content"])
            metadatas.append(doc.get("metadata", {}))
        embeddings_list = embeddings.embed_documents(texts)
        self.collection.add(ids=ids, documents=texts, embeddings=embeddings_list, metadatas=metadatas)

    def similarity_search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        results = self.collection.query(
            query_texts=[query],
            n_results=k,
            include=['documents', 'distances', 'metadatas'],
        )
        docs = []
        for doc, dist, meta in zip(results["documents"][0], results["distances"][0], results["metadatas"][0]):
            docs.append({"content": doc, "distance": dist, "metadata": meta})
        return docs

# Singleton instance
vector_store = ChromaVectorStore(col)