"""
Main agent logic: decides whether to use local KB or web search.
"""

import os
from typing import Dict, Any

from langchain_ollama import ChatOllama
from langchain.agents import initialize_agent, AgentType, Tool, AgentExecutor
from langchain_core.messages import HumanMessage

from rag_tools import search_local_kb, web_search
from vectorstore import create_vectorstore, load_documents

# Load or create vector store
vectorstore = create_vectorstore()
# Load documents from the documents folder if not already loaded
if not vectorstore._collection.count():  # type: ignore[attr-defined]
    load_documents("./documents", vectorstore)

# Define tools
tools = [
    Tool(name="search_local_kb", func=search_local_kb, description="Search the local knowledge base."),
    Tool(name="web_search", func=web_search, description="Search the web using Tavily."),
]

# System prompt to guide the agent
system_prompt = (
    "You are an AI assistant. For questions about local documents use the 'search_local_kb' tool. "
    "For recent news or facts not in the local docs, use 'web_search'. "
    "Always indicate the source of your answer (chromadb or tavily)."
)

# Create the agent executor
llm = ChatOllama(model="llama3")
agent_executor = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.OPENAI_FUNCTIONS,
    verbose=True,
    system_message=system_prompt,
)

def main():
    print("Welcome to the RAG agent. Type 'exit' to quit.")
    while True:
        user_input = input("\nUser: ")
        if user_input.lower() in {"exit", "quit"}:
            print("Goodbye!")
            break
        # Run the agent
        result = agent_executor.invoke({"input": user_input})
        # The result may contain tool calls and final answer
        print("\nAssistant:", result.get("output", ""))

if __name__ == "__main__":
    main()