"""Module for interacting with Qdrant vector store using Ollama embeddings.
"""

from typing import List, Dict, Any

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient

class QdrantStore:
    """Wrapper around QdrantVectorStore.

    Parameters
    ----------
    collection_name: str
        Name of the Qdrant collection.
    host: str
        Qdrant host URL.
    port: int
        Qdrant port.
    """

    def __init__(self, collection_name: str = "rag_collection", host: str = "localhost", port: int = 6333):
        self.collection_name = collection_name
        self.client = QdrantClient(host=host, port=port)
        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
        # Create collection if not exists
        if collection_name not in self.client.get_collections().collections:
            self.client.recreate_collection(collection_name=collection_name, vectors_config={"size": 512, "distance": "Cosine"})
        self.store = QdrantVectorStore.from_existing_collection(
            collection_name=collection_name,
            embedding=self.embeddings,
            client=self.client,
        )

    def add_documents(self, documents: List[str], titles: List[str] | None = None, metadatas: List[Dict[str, Any]] | None = None) -> None:
        """Add documents to the collection.

        Each document is added as a separate vector. If titles or metadatas are provided, they are attached.
        """
        if titles is None:
            titles = [f"doc_{i}" for i in range(len(documents))]
        if metadatas is None:
            metadatas = [{} for _ in range(len(documents))]
        self.store.add_texts(texts=documents, metadatas=metadatas, ids=titles)

    def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
        """Semantic search returning list of dicts with 'content' and 'metadata'."""
        results = self.store.similarity_search_with_score(query, k=k)
        return [{"content": r[0].page_content, "metadata": r[0].metadata, "score": r[1]} for r in results]