"""Core logic for the RAG agent.

The agent decides whether to use the local knowledge base or Tavily based on a
very simple heuristic: if the query contains words like ``news``, ``latest``
or ``today`` it is routed to the web search; otherwise the local KB is used.

The decision logic can be replaced with a more sophisticated router if
desired.
"""

from __future__ import annotations

from typing import Tuple

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

# ---------------------------------------------------------------------------
# Prompt template
# ---------------------------------------------------------------------------
SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search.

When answering, always include the source of the information:
- "chromadb" for local knowledge base results.
- "tavily" for web search results.

If you are uncertain, say "I don't know" but still mention the source you used.
"""

USER_PROMPT = """Question: {question}\n
When you respond, first state the source (chromadb or tavily) and then provide the answer.
"""

prompt = ChatPromptTemplate.from_messages([
    ("system", SYSTEM_PROMPT),
    ("user", USER_PROMPT),
])

# ---------------------------------------------------------------------------
# Decision logic
# ---------------------------------------------------------------------------
WEB_KEYWORDS = {"news", "latest", "today", "current", "recent"}


def choose_tool(question: str) -> Tuple[str, callable]:
    """Return the name of the tool and the function to call.

    Parameters
    ----------
    question: str
        The user query.

    Returns
    -------
    Tuple[str, callable]
        The tool name and the corresponding function.
    """
    lowered = question.lower()
    if any(word in lowered for word in WEB_KEYWORDS):
        return "web_search", web_search
    return "search_local_kb", search_local_kb

# ---------------------------------------------------------------------------
# Agent creation
# ---------------------------------------------------------------------------

def create_agent() -> AgentExecutor:
    """Instantiate the agent with the two tools.

    The LLM used is Ollama's ``llama3``.
    """
    tools = [search_local_kb, web_search]
    llm = ChatOllama(model="llama3", temperature=0)
    # Build an agent that knows about the tools and uses the custom prompt
    agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt)
    return AgentExecutor(agent=agent, tools=tools, verbose=True)

# ---------------------------------------------------------------------------
# CLI loop
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # Ensure the vector store is loaded once
    store = create_vectorstore()
    # Load documents if the store is empty
    if not store.get_index_info():
        from vectorstore import load_documents
        load_documents("documents", store)

    agent = create_agent()
    print("RAG Agent ready. Type 'exit' to quit.")
    while True:
        try:
            question = input("\nЗапрос: ")
        except EOFError:
            break
        if question.strip().lower() in {"exit", "quit"}:
            break
        # The agent will automatically call the chosen tool via the prompt.
        # We simply pass the question to the agent.
        result = agent.invoke({"input": question})
        # The agent's output already contains the source.
        print("Ответ:", result["output"])  # noqa: T201