"""Main script for the RAG agent with ChromaDB and Tavily.

The script:
1. Loads or creates the Chroma vector store.
2. Loads documents from the `documents/` folder.
3. Sets up the LangChain agent with two tools: `search_local_kb` and `web_search`.
4. Runs a simple CLI loop.
"""

import os
from pathlib import Path

from langchain_ollama import ChatOllama
from langchain.agents import initialize_agent, AgentType
from langchain.tools import Tool

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

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

# ---------------------------------------------------------------------------
# Helper: load or create vector store
# ---------------------------------------------------------------------------
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)

# Load documents – we always load; Chroma will deduplicate by ID if same content
print("Loading documents into ChromaDB (if not already present)...")
load_documents(DOCS_DIR, vectorstore)
print("Documents loaded.")

# ---------------------------------------------------------------------------
# Define tools – pass the vectorstore to the local search tool
# ---------------------------------------------------------------------------
# We wrap the tool functions to include the vectorstore argument

def local_kb_tool(query: str, top_k: int = 3):
    return search_local_kb(query=query, top_k=top_k, vectorstore=vectorstore)

# Create LangChain Tool objects
local_tool = Tool(
    name="search_local_kb",
    func=local_kb_tool,
    description="Semantic search in the local knowledge base. Use when the answer is in the local documents.",
)
web_tool = Tool(
    name="web_search",
    func=web_search,
    description="Search the web using Tavily. Use for up‑to‑date facts or news.",
)

# ---------------------------------------------------------------------------
# Agent setup
# ---------------------------------------------------------------------------
llm = ChatOllama(model=MODEL, temperature=0.0)

system_prompt = (
    "You are an assistant that answers user questions. "
    "If the answer is likely to be in the local knowledge base, use the tool "
    "search_local_kb. If the answer requires up‑to‑date information, use the "
    "web_search tool. After retrieving information, provide the answer and "
    "state the source: either 'chromadb' or 'tavily'."
)

agent = initialize_agent(
    tools=[local_tool, web_tool],
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    prefix=system_prompt,
)

# ---------------------------------------------------------------------------
# CLI loop
# ---------------------------------------------------------------------------
print("\nRAG Agent ready. Type your question (or 'exit' to quit).\n")
while True:
    try:
        query = input("Query: ")
    except (KeyboardInterrupt, EOFError):
        print("\nExiting.")
        break
    if query.strip().lower() in {"exit", "quit", "q"}:
        print("Exiting.")
        break
    if not query.strip():
        continue
    # Run the agent
    try:
        result = agent.run(query)
        print(f"\nAnswer:\n{result}\n")
    except Exception as e:
        print(f"Error: {e}")
        continue