"""Core logic for the RAG agent.

The agent decides whether to use the local knowledge base or Tavily based on simple heuristics:

- If the query contains words that usually refer to recent events (e.g. "новости", "актуальные", "сегодня", "сейчас"), we route to Tavily.
- Otherwise we assume the answer can be found in the local KB.

The agent uses LangChain's :class:`langchain.agents.AgentExecutor` with a custom prompt that
instructs the LLM to specify the source in the response.
"""

import os
from typing import Dict, Any

from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain_ollama import ChatOllama
from langchain_core.prompts import ChatPromptTemplate

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

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
# Directory containing the documents to index
DOCUMENTS_DIR = "./documents"
# Persistence directory for Chroma
CHROMA_DIR = "./chroma_db"
# The LLM used by the agent
LLM = ChatOllama(model="llama3")

# ---------------------------------------------------------------------------
# Helper functions
# ---------------------------------------------------------------------------

def should_use_web(query: str) -> bool:
    """Return True if the query looks like it needs up‑to‑date information.

    The heuristic looks for a few Russian keywords that usually indicate a
    request for recent news.
    """
    keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
    lowered = query.lower()
    return any(k in lowered for k in keywords)

# ---------------------------------------------------------------------------
# Agent setup
# ---------------------------------------------------------------------------
# Load or create the vector store
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
# If the store is empty, load documents from the documents directory
if not vectorstore._collection.count():  # type: ignore[attr-defined]
    load_documents(DOCUMENTS_DIR, vectorstore)

# Tools – we pass the vectorstore instance to the local search tool via a closure
local_kb_tool = search_local_kb
local_kb_tool.__globals__["vectorstore"] = vectorstore  # inject vectorstore

TOOLS = [local_kb_tool, web_search]

# Prompt template – the LLM is instructed to specify the source.
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are an AI assistant that answers user questions.") ,
    ("user", "{input}"),
])

# Agent – we use the simple tool‑calling agent
agent = create_openai_tools_agent(llm=LLM, tools=TOOLS, prompt=prompt)
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)

# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    print("RAG Agent ready. Type 'exit' to quit.")
    while True:
        user_query = input("\nЗапрос: ")
        if user_query.lower() in {"exit", "quit", "q"}:
            break
        # Decide which tool to use
        if should_use_web(user_query):
            # Explicitly call the web search tool
            result = web_search(user_query)
            source = "tavily"
        else:
            result = search_local_kb(user_query, vectorstore=vectorstore)
            source = "chromadb"
        # Ask the LLM to format the answer
        formatted = LLM.invoke({"input": user_query + "\n\nAnswer: " + result})
        print("\nОтвет:", formatted.content)
        print("Источник:", source)

if __name__ == "__main__":
    main()

# End of agent.py