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

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

from langchain_ollama import OllamaEmbeddings, Ollama
from langchain_qdrant import QdrantVectorStore
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.tools import tool
from langchain.agents import create_agent, AgentExecutor, AgentType
from langchain.schema import AgentAction, AgentFinish
from langchain.callbacks import get_openai_callback

# Configuration
QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_agent")
Ollama_EMBED_MODEL = os.getenv("Ollama_EMBED_MODEL", "nomic-embed-text")
Ollama_MODEL = os.getenv("Ollama_MODEL", "llama3")

# Initialize embeddings and vector store
embeddings = OllamaEmbeddings(model=Ollama_EMBED_MODEL)
vector_store = QdrantVectorStore(
    url=QDRANT_URL,
    collection_name=QDRANT_COLLECTION,
    embeddings=embeddings,
)

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

# Tools
@tool
def search_knowledge_base(query: str, max_results: int = 5) -> str:
    """Semantic search in the knowledge base."""
    docs = vector_store.similarity_search(query, k=max_results)
    if not docs:
        return "No relevant documents found."
    return "\n\n".join([f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)])

@tool
def add_to_knowledge_base(content: str, title: str) -> str:
    """Add a document to the knowledge base."""
    # Split content into chunks
    chunks = splitter.split_text(content)
    # Create documents with metadata
    docs = []
    for i, chunk in enumerate(chunks):
        docs.append(
            {
                "page_content": chunk,
                "metadata": {"title": title, "chunk_index": i},
            }
        )
    # Add to vector store
    vector_store.add_documents(docs)
    return f"Document '{title}' added with {len(chunks)} chunks."

# Agent setup
SYSTEM_PROMPT = (
    "You are an AI assistant with access to a local knowledge base. "
    "Use the provided tools to search and add information. "
    "When answering, rely on the knowledge base and the LLM."
)

# Create agent with tools
agent = create_agent(
    llm=Ollama(model=Ollama_MODEL),
    tools=[search_knowledge_base, add_to_knowledge_base],
    system_prompt=SYSTEM_PROMPT,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
)

executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base], verbose=True)

# CLI
def main():
    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /quit")
    while True:
        try:
            inp = input("> ")
        except EOFError:
            break
        if not inp:
            continue
        if inp.startswith("/quit"):
            print("Bye!")
            break
        if inp.startswith("/add"):
            parts = inp.split(maxsplit=2)
            if len(parts) < 3:
                print("Usage: /add <title> <content>")
                continue
            title, content = parts[1], parts[2]
            print(add_to_knowledge_base(content, title))
            continue
        if inp.startswith("/search"):
            query = inp[len("/search"):].strip()
            if not query:
                print("Usage: /search <query>")
                continue
            print(search_knowledge_base(query))
            continue
        # Treat as normal user query
        response = executor.invoke({"input": inp})
        print(response.get("output", ""))

if __name__ == "__main__":
    main()