from pathlib import Path
from uuid import uuid4
from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document

class KnowledgeBase:
    """A simple wrapper around Qdrant for storing and searching documents."""
    def __init__(self, collection_name: str = "rag_kb", persist_path: str = "qdrant_data"):
        # Initialize Qdrant client (in‑memory or on‑disk)
        self.client = QdrantClient(path=persist_path)
        # Create collection with cosine distance and 3072‑dim vectors (nomic‑embed‑text output)
        self.client.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=3072, distance=Distance.COSINE),
        )
        # Create vector store wrapper
        self.vector_store = QdrantVectorStore(
            client=self.client,
            collection_name=collection_name,
            embedding=OllamaEmbeddings(model="nomic-embed-text"),
        )
        # Text splitter for chunking documents
        self.splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

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

        The content is split into chunks, each chunk is embedded and stored.
        """
        # Create Document objects with metadata
        docs = self.splitter.create_documents([content], metadata={"title": title})
        # Add to vector store
        self.vector_store.add_documents(docs)

    def search(self, query: str, limit: int = 10):
        """Search the knowledge base for the most relevant documents.

        Returns a list of dictionaries containing page_content, metadata and score.
        """
        results = self.vector_store.similarity_search_with_score(query, k=limit)
        return [
            {
                "page_content": doc.page_content,
                "metadata": doc.metadata,
                "score": score,
            }
            for doc, score in results
        ]

# Global singleton instance used by tools and the agent
kb = KnowledgeBase()