"""Vector store implementation using Qdrant and Ollama embeddings.

This module provides a KnowledgeBase class that wraps a QdrantVectorStore and exposes
methods for adding documents and searching the knowledge base.
"""

from pathlib import Path
from typing import List, Dict, Any

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from qdrant_client import QdrantClient

class KnowledgeBase:
    """A simple wrapper around QdrantVectorStore.

    Parameters
    ----------
    collection_name: str
        Name of the Qdrant collection to use.
    host: str, optional
        Qdrant host address. Defaults to ``localhost``.
    port: int, optional
        Qdrant port. Defaults to ``6333``.
    """

    def __init__(self, collection_name: str = "knowledge_base", host: str = "localhost", port: int = 6333):
        self.collection_name = collection_name
        self.client = QdrantClient(host=host, port=port)
        # Ensure the collection exists
        self.client.recreate_collection(
            collection_name=self.collection_name,
            vectors_config={"size": 1024, "distance": "Cosine"},
        )

        self.embeddings = OllamaEmbeddings(model="nomic-embed-text")
        self.vector_store = QdrantVectorStore(
            client=self.client,
            collection_name=self.collection_name,
            embeddings=self.embeddings,
        )
        self.splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

    def add_document(self, content: str, title: str) -> None:
        """Add a document to the knowledge base.

        The document is split into chunks, embedded, and stored in Qdrant.
        """
        chunks = self.splitter.split_text(content)
        documents: List[Document] = []
        for idx, chunk in enumerate(chunks):
            meta = {"title": title, "chunk_index": idx}
            documents.append(Document(page_content=chunk, metadata=meta))
        self.vector_store.add_documents(documents)

    def search(self, query: str, max_results: int = 5) -> List[Dict[str, Any]]:
        """Search the knowledge base for the most relevant chunks.

        Returns a list of dictionaries containing the chunk content, title, and score.
        """
        results = self.vector_store.similarity_search_with_score(query, k=max_results)
        output = []
        for doc, score in results:
            output.append({
                "content": doc.page_content,
                "title": doc.metadata.get("title"),
                "chunk_index": doc.metadata.get("chunk_index"),
                "score": score,
            })
        return output

    def load_from_directory(self, directory: str) -> None:
        """Load all .txt files from a directory into the knowledge base.

        Parameters
        ----------
        directory: str
            Path to the directory containing text files.
        """
        path = Path(directory)
        for file_path in path.rglob("*.txt"):
            title = file_path.stem
            content = file_path.read_text(encoding="utf-8")
            self.add_document(content, title)

# Example usage (uncomment for quick test)
# if __name__ == "__main__":
#     kb = KnowledgeBase()
#     kb.load_from_directory("data")
#     print(kb.search("What is Python?", 3))