import os

from langchain.agents import AgentExecutor, create_structured_chat_agent
from langchain.memory import ConversationBufferMemory
from langchain_core.messages import AIMessage
from langchain_core.runnables.config import RunnableConfig
from langchain_openai import AzureChatOpenAI

from src import tools
from src.prompts import solver_prompt
from src.state import AgentState

# This node uses the REACT framework to solve subqueries
tools_list = [getattr(tools, name) for name in tools.__all__]


async def solver_node(state: AgentState):
    """Executes a single step of the agent's problem-solving process using
    a ReAct agent.

    This function retrieves the current subquery from the agent's state,
    constructs a prompt including the context of previous results, and invokes
    the agent to generate a response. The agent's output is then appended to
    the results, the current step is incremented, and the new message is added
    to the conversation history.

    Args:
        state (AgentState): The current state of the agent, containing
                            subqueries, results, current step index, and
                            message history.
    Returns:
        dict: An updated state dictionary with the new result, incremented
              step, and updated messages.
    """
    # Create the agent with the REACT framework
    agent = create_structured_chat_agent(
        AzureChatOpenAI(deployment_name=os.getenv(
            "AZURE_OPENAI_DEPLOYMENT_NAME")),
        tools_list,
        solver_prompt
    )

    # Set up memory to retain conversation context
    memory = ConversationBufferMemory(
        memory_key="chat_history", return_messages=True)

    # Create the agent executor
    agent_executor = AgentExecutor(
        agent=agent, tools=tools_list, memory=memory,
        verbose=True, handle_parsing_errors=True)

    # Set the recursion limit for the agent
    config = RunnableConfig(
        recursion_limit=32
    )

    # Get the current subquery
    current_subquery = state["subqueries"][state["current_step"]]

    # Add previous results to context
    context = "\n".join(
        [f"Step {i+1}: {res}" for i, res in enumerate(state["results"])])
    full_query = (
        f"Original Query: {state['original_query']}\n\n"
        f"Context:\n{context}\n\n"
        f"Current Task: {current_subquery}"
    )

    # Invoke the agent with the full query
    response = await agent_executor.ainvoke(
        {"input": full_query}, config)
    output = response['output']

    return {
        **state,
        "results": state["results"] + [output],
        "current_step": state["current_step"] + 1,
        "messages": state["messages"] + [AIMessage(content=output)]
    }