"""
Agent with RAG memory using Qdrant and Ollama.

This is a minimal example that demonstrates:
1. Connecting to a local Qdrant instance.
2. Creating two tools: `search_knowledge_base` and `add_to_knowledge_base`.
3. Using LangChain's RecursiveCharacterTextSplitter to chunk documents.
4. Building an agent with the tools via `create_agent`.
5. A simple REPL that accepts `/add`, `/search` and `/quit` commands.

To run:
    pip install -r requirements.txt
    python main.py

Make sure a Qdrant instance is running locally (default port 6333) and Ollama is available at http://localhost:11434.
"""

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

from langchain_community.document_loaders import DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.embeddings.ollama import OllamaEmbeddings
from langchain.vectorstores.qdrant import QdrantVectorStore
from langchain.agents import create_agent, AgentExecutor, Tool
from langchain.schema import Document

# Configuration
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
OLLAMA_MODEL = os.getenv("OLLAMA_MODEL", "llama2")
COLLECTION_NAME = "rag_collection"

# Initialize embeddings and vector store
embeddings = OllamaEmbeddings(model=OLLAMA_MODEL)
vector_store = QdrantVectorStore(
    url=QDRANT_URL,
    collection_name=COLLECTION_NAME,
    embedding_function=embeddings,
)

# Tool: search knowledge base

def search_knowledge_base(query: str) -> str:
    """Return the top 3 relevant snippets for a query."""
    docs = vector_store.similarity_search_with_score(query, k=3)
    if not docs:
        return "No relevant documents found."
    results = [f"{idx+1}. {doc[0].page_content[:200]}… (score: {doc[1]:.4f})" for idx, doc in enumerate(docs)]
    return "\n".join(results)

# Tool: add to knowledge base

def add_to_knowledge_base(file_path: str) -> str:
    """Load a text file, chunk it and add to Qdrant."""
    loader = DirectoryLoader(Path(file_path).parent.as_posix(), glob=Path(file_path).name)
    docs = loader.load()
    if not docs:
        return f"No documents found in {file_path}."
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    chunks: List[Document] = []
    for doc in docs:
        chunks.extend(splitter.split_documents([doc]))
    vector_store.add_documents(chunks)
    return f"Added {len(chunks)} chunks from {file_path} to the knowledge base."

# Define tools list
TOOLS: List[Tool] = [
    Tool(
        name="search_knowledge_base",
        func=search_knowledge_base,
        description="Search the local Qdrant knowledge base for relevant information.",
    ),
    Tool(
        name="add_to_knowledge_base",
        func=add_to_knowledge_base,
        description="Add a text file to the knowledge base. Provide full path.",
    ),
]

# Build agent
agent = create_agent(TOOLS, llm=embeddings)  # embeddings can act as LLM via Ollama
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)

# Simple REPL
if __name__ == "__main__":
    print("Welcome to the RAG agent. Commands: /add <file>, /search <query>, /quit")
    while True:
        try:
            inp = input(">>> ")
        except EOFError:
            break
        if not inp:
            continue
        if inp.startswith("/quit"):
            print("Goodbye!")
            break
        elif inp.startswith("/add "):
            path = inp.split(maxsplit=1)[1]
            print(add_to_knowledge_base(path))
        elif inp.startswith("/search "):
            query = inp.split(maxsplit=1)[1]
            print(search_knowledge_base(query))
        else:
            # Treat as normal agent prompt
            result = executor.invoke({"input": inp})
            print(result.get("output", ""))