"""Agent construction for RAG with local KB and web search.

The agent uses LangChain's `create_openai_functions_agent` style with a custom
`AgentExecutor` that routes queries to either the local knowledge base or the
web search based on a simple heuristic: if the query contains words that
suggest recent events ("news", "today", "now", "latest"), we use the web
search; otherwise we use the local KB.

The agent returns the answer along with the source used.
"""

from typing import Dict, Any

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

from rag_tools import search_local_kb, web_search
from vectorstore import create_vectorstore

# ---------------------------------------------------------------------------
# Helper: decide whether to use web search
# ---------------------------------------------------------------------------

WEB_KEYWORDS = {"news", "today", "now", "latest", "recent", "current"}


def should_use_web(query: str) -> bool:
    words = set(query.lower().split())
    return bool(words & WEB_KEYWORDS)

# ---------------------------------------------------------------------------
# System prompt
# ---------------------------------------------------------------------------

SYSTEM_PROMPT = "You are an assistant that answers questions. Use the local knowledge base when the question is about stored documents; otherwise use web search. Respond with the answer and indicate the source (chromadb or tavily)."

# ---------------------------------------------------------------------------
# Agent construction
# ---------------------------------------------------------------------------

def create_agent(vectorstore) -> AgentExecutor:
    """Create an AgentExecutor that routes queries to the appropriate tool.

    Parameters
    ----------
    vectorstore: Chroma
        The vector store used by the local KB search tool.

    Returns
    -------
    AgentExecutor
        Configured agent.
    """
    # Attach the vectorstore to the local search tool via closure
    def search_local_kb_wrapper(query: str, top_k: int = 3) -> str:
        return search_local_kb(query, top_k=top_k, vectorstore=vectorstore)

    # Wrap the web search tool (no extra params needed)
    def web_search_wrapper(query: str, top_k: int = 3) -> str:
        return web_search(query, top_k=top_k)

    # Define tools list with updated references
    tools = [search_local_kb_wrapper, web_search_wrapper]

    # LLM
    llm = ChatOllama(model="llama3")

    # Prompt template
    prompt = ChatPromptTemplate.from_messages([
        ("system", SYSTEM_PROMPT),
        ("human", "{input}"),
    ])

    # Create agent
    agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)

    # Wrap with AgentExecutor
    return AgentExecutor(agent=agent, tools=tools, verbose=True)

# ---------------------------------------------------------------------------
# Simple executor for tests
# ---------------------------------------------------------------------------

def create_agent_executor(vectorstore) -> Any:
    """Return a simple callable that mimics the agent for testing.

    The returned callable takes a dictionary with keys ``input`` and optionally
    ``tool_choice``. It selects the appropriate tool based on the query and
    returns a dictionary with an ``output`` key containing the result.
    """
    agent = create_agent(vectorstore)

    def executor(inputs: Dict[str, Any]) -> Dict[str, Any]:
        query = inputs.get("input", "")
        tool_choice = inputs.get("tool_choice")
        if tool_choice == "web_search":
            result = web_search(query)
        elif tool_choice == "search_local_kb":
            result = search_local_kb(query, vectorstore=vectorstore)
        else:
            # Default routing
            if should_use_web(query):
                result = web_search(query)
            else:
                result = search_local_kb(query, vectorstore=vectorstore)
        return {"output": result}

    return executor

# ---------------------------------------------------------------------------
# Example usage
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    store = create_vectorstore()
    agent = create_agent(store)
    while True:
        q = input("Query> ")
        if q.lower() in {"exit", "quit", "q"}:
            break
        # Simple routing: if query contains web keywords, use web tool
        if should_use_web(q):
            result = agent.invoke({"input": q, "tool_choice": "web_search"})
        else:
            result = agent.invoke({"input": q, "tool_choice": "search_local_kb"})
        print("Answer:", result.get("output", ""))
""