"""RAG agent CLI.

This script demonstrates a simple chat loop with a LangChain agent that
searches either a local Chroma vector store or the web via Tavily.  The
agent automatically decides which tool to use based on the user query.

Prerequisites:
  * Ollama must be running locally with the ``llama3`` model and the
    ``nomic-embed-text`` embedding model.
  * A valid Tavily API key must be set in the environment variable
    ``TAVILY_API_KEY``.
  * The ``documents`` directory should contain the source text files.
"""

import os
from pathlib import Path

from langchain_ollama import ChatOllama
from langchain.agents import create_agent

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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
CHROMA_DIR = "./chroma_db"
DOCS_DIR = "./documents"

# ---------------------------------------------------------------------------
# Initialise vector store and load documents
# ---------------------------------------------------------------------------
print("Initializing Chroma vector store…")
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
print("Loading documents…")
load_documents(DOCS_DIR, vectorstore)

# ---------------------------------------------------------------------------
# Agent setup
# ---------------------------------------------------------------------------
# System prompt that tells the model how to choose a tool.
SYSTEM_PROMPT = (
    "You are an AI assistant that can answer questions using two tools. "
    "If the answer requires up‑to‑date information, use the web_search tool. "
    "Otherwise, use the search_local_kb tool. "
    "When you call a tool, the tool will return the answer. "
    "Respond with the final answer and include the source tag (chromadb or tavily)."
)

llm = ChatOllama(model="llama3", temperature=0)

# Tools list
TOOLS = [search_local_kb, web_search]

agent = create_agent(
    model=llm,
    tools=TOOLS,
    system_prompt=SYSTEM_PROMPT,
)

# ---------------------------------------------------------------------------
# Chat loop
# ---------------------------------------------------------------------------
print("\n--- RAG Agent CLI ---")
print("Type 'exit' or 'quit' to end.")
while True:
    user_input = input("\nUser: ")
    if user_input.lower() in {"exit", "quit", "q"}:
        print("Goodbye!")
        break
    # Invoke the agent
    try:
        response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
        # The response is a dict with a "messages" key
        assistant_msg = next(
            m for m in response["messages"] if m["role"] == "assistant"
        )
        print("\nAssistant:", assistant_msg["content"].strip())
    except Exception as e:
        print("Error:", e)

"""End of main.py"""