import os
from typing import Any, Dict

from langchain_ollama import Ollama
from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
from langchain.schema import AgentAction, AgentFinish
from langchain.tools import tool
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.output_parsers import StrOutputParser

from vectorstore import QdrantVectorStore
from tools import search_local_kb, web_search


def create_agent(
    vectorstore: QdrantVectorStore,
    tavily_api_key: str,
    llm_model: str = "llama3",
) -> AgentExecutor:
    """
    Create a LangChain agent that routes queries to either the local KB or the web.

    Parameters
    ----------
    vectorstore : QdrantVectorStore
        The vector store for local knowledge base.
    tavily_api_key : str
        Tavily API key.
    llm_model : str
        Ollama model name for LLM.

    Returns
    -------
    AgentExecutor
        Configured agent executor.
    """
    # LLM
    llm = Ollama(model=llm_model)

    # Define tools with partial application of required arguments
    tools = [
        Tool(
            name="search_local_kb",
            func=lambda query, top_k=3: search_local_kb(
                query=query, top_k=top_k, vectorstore=vectorstore
            ),
            description=(
                "Use this tool to search the local knowledge base. "
                "Return the most relevant snippets."
            ),
        ),
        Tool(
            name="web_search",
            func=lambda query, max_results=3: web_search(
                query=query, tavily_api_key=tavily_api_key, max_results=max_results
            ),
            description=(
                "Use this tool to search the web via Tavily. "
                "Return the most relevant snippets."
            ),
        ),
    ]

    # Prompt template for the agent
    system_prompt = (
        "You are an AI assistant that can answer questions using either "
        "the local knowledge base or the web. If the question is about "
        "recent events, news, or requires up-to-date information, "
        "use the web_search tool. If the question is about "
        "information contained in the local documents, use the "
        "search_local_kb tool. After retrieving the information, "
        "provide a concise answer and state the source (chromadb or tavily)."
    )

    prompt = ChatPromptTemplate.from_messages(
        [
            ("system", system_prompt),
            MessagesPlaceholder("history"),
            ("human", "{input}"),
            MessagesPlaceholder("agent_scratchpad"),
        ]
    )

    # Output parser
    output_parser = StrOutputParser()

    # Agent
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
        verbose=False,
        handle_parsing_errors=True,
        max_iterations=5,
        early_stopping_method="generate",
        system_message=system_prompt,
    )

    return agent