import os

from langchain_core.messages import HumanMessage
from langchain_core.output_parsers import JsonOutputParser
from langchain_openai import AzureChatOpenAI

from src.prompts import simplifier_prompt
from src.state import AgentState


async def simplifier_node(state: AgentState):
    """Processes the agent's state to generate simplified subqueries
    from the original query using a language model.

    This function takes the current agent state, sends the original query to
    a language model via a prompt chain, and parses the output to extract a
    list of simpler/atomic subqueries. It then updates the state with the
    generated subqueries, resets the current step and results, and appends
    the original query as a message.

    Args:
        state (AgentState): The current state of the agent, expected to
                            contain the key "original_query".
    Returns:
        dict: An updated state dictionary containing:
            - "subqueries": List of subqueries generated from the
                            original query.
            - "current_step": Reset to 0.
            - "results": Reset to an empty list.
            - "messages": List containing the original query as a HumanMessage.
            - All other original state keys and values.
    """
    # Create a message with the original query
    messages = [HumanMessage(content=state["original_query"])]

    # Create the simplifier chain
    chain = (
        simplifier_prompt |
        AzureChatOpenAI(deployment_name=os.getenv(
            "AZURE_OPENAI_DEPLOYMENT_NAME")) |
        JsonOutputParser()
    )

    # Invoke the simplifier chain
    result = await chain.ainvoke({"query": state["original_query"]})

    return {
        **state,
        "subqueries": result["subqueries"],
        "messages": messages
    }