"""Tools for RAG agent.

This module defines two LangChain tools:

* ``search_knowledge_base`` – semantic search in Qdrant.
* ``add_to_knowledge_base`` – add a document to Qdrant.

The tools use the global ``vector_store`` instance defined in this module.
"""

from typing import List

from langchain_ollama import OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.schema import Document

# Global vector store instance
# We initialise it lazily – the first call to the tools will create the store.
_vector_store = None

# Embedding model
_embedding = OllamaEmbeddings(model="nomic-embed-text")

# Text splitter
_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)


def _get_vector_store() -> QdrantVectorStore:
    global _vector_store
    if _vector_store is None:
        # Create a new collection named "knowledge_base"
        _vector_store = QdrantVectorStore.from_texts(
            texts=[],
            embedding=_embedding,
            url="http://localhost:6333",
            collection_name="knowledge_base",
        )
    return _vector_store


@tool("search_knowledge_base")
def search_knowledge_base(query: str, max_results: int = 5) -> List[dict]:
    """Semantic search in the knowledge base.

    Parameters
    ----------
    query: str
        The search query.
    max_results: int
        Maximum number of results to return.

    Returns
    -------
    List[dict]
        List of matching documents with ``content`` and ``metadata``.
    """
    store = _get_vector_store()
    results = store.similarity_search_with_score(query, k=max_results)
    # Return only content and metadata for simplicity
    return [
        {"content": doc.page_content, "metadata": doc.metadata}
        for doc, _ in results
    ]


@tool("add_to_knowledge_base")
def add_to_knowledge_base(content: str, title: str) -> str:
    """Add a document to the knowledge base.

    The content is split into chunks before being added.
    """
    store = _get_vector_store()
    chunks = _splitter.split_text(content)
    docs: List[Document] = []
    for i, chunk in enumerate(chunks):
        docs.append(
            Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
        )
    store.add_documents(docs)
    return f"Document '{title}' added with {len(chunks)} chunks."