"""
Main entry point for the RAG agent.
"""

import asyncio
import os
from pathlib import Path

from langchain_ollama import ChatOllama
from langchain.agents import create_agent
from langchain.agents.agent_toolkits import BaseToolkit
from langchain.agents.agent_types import AgentType
from langchain.tools import BaseTool

from rag_tools import search_knowledge_base, add_to_knowledge_base
from qdrant_store import load_directory

# Load environment variables if any
from dotenv import load_dotenv
load_dotenv()

# Configuration
LLM_MODEL = "llama3"
KNOWLEDGE_DIR = os.getenv("KNOWLEDGE_DIR", "./knowledge")

# Ensure knowledge directory exists and load documents
Path(KNOWLEDGE_DIR).mkdir(parents=True, exist_ok=True)
load_directory(KNOWLEDGE_DIR)

# Define tools
class SearchTool(BaseTool):
    name = "search_knowledge_base"
    description = "Perform semantic search in the knowledge base."
    func = search_knowledge_base

class AddTool(BaseTool):
    name = "add_to_knowledge_base"
    description = "Add a new document to the knowledge base."
    func = add_to_knowledge_base

# Simple toolkit
class RAGToolkit(BaseToolkit):
    def get_tools(self):
        return [SearchTool(), AddTool()]

    def get_base_prompt(self):
        return None

# Create LLM
llm = ChatOllama(model=LLM_MODEL)

# System prompt instructing the agent to use the knowledge base
SYSTEM_PROMPT = """
You are an assistant that uses a knowledge base. When answering user queries, first search the knowledge base with the search_knowledge_base tool. If the information is not sufficient, ask the user for clarification. You can also add new documents to the knowledge base using add_to_knowledge_base.
"""

# Create agent
agent = create_agent(
    llm=llm,
    toolkit=RAGToolkit(),
    system_prompt=SYSTEM_PROMPT,
    agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
)

async def main():
    print("RAG Agent ready. Type your query (or 'quit' to exit).")
    while True:
        user_input = input("\n> ")
        if user_input.lower() in {"quit", "exit", "q"}:
            print("Goodbye!")
            break
        response = await agent.ainvoke(user_input)
        print("\nAssistant:", response)

if __name__ == "__main__":
    asyncio.run(main())