from langgraph.graph import END, StateGraph

from src.nodes import refiner_node, simplifier_node, solver_node

import asyncio
from src.state import AgentState


def should_continue(state: AgentState) -> str:
    """Determines whether to continue solving subqueries or
    refine the final answer.

    Args:
        state (AgentState): The current state of the agent, containing
                            subqueries, results, current step index, and
                            message history.

    Returns:
        str: A string indicating the next step:
             - "continue_solve": Continue solving subqueries.
             - "refine": Refine the final answer.
    """
    # Check if we have any subqueries to process
    if not state["subqueries"]:
        return "refine"

    current_step = state.get("current_step", 0)

    # Check if we've processed all subqueries
    if current_step >= len(state["subqueries"]):
        return "refine"

    # Continue processing
    return "continue_solve"


class AgentGAIA:
    """Agent GAIA for multi-step query resolution."""

    def __init__(self):
        """Initialize the Agent GAIA with a state graph."""
        # Create the state graph
        builder = StateGraph(AgentState)

        # Add nodes to the graph
        builder.add_node("simplify", simplifier_node)
        builder.add_node("solve", solver_node)
        builder.add_node("refine", refiner_node)

        # Set the initial state
        builder.set_entry_point("simplify")

        # Define transitions
        builder.add_edge("simplify", "solve")

        builder.add_conditional_edges(
            "solve",
            should_continue,
            {
                "continue_solve": "solve",
                "refine": "refine"
            }
        )

        builder.add_edge("refine", END)

        self.graph = builder.compile()

    async def __call__(self, query: str) -> str:
        """Run a query through the GAIA Agent.
        Args:
            query (str): The initial query to be processed.
        Returns:
            str: The final answer generated by the agent.
        """
        state = {
            "original_query": query,
            "subqueries": [],
            "current_step": 0,
            "results": [],
            "final_answer": None,
            "messages": []
        }

        state = await self.graph.ainvoke(state)

        return state["final_answer"] or state["results"][-1]


if __name__ == "__main__":
    agent = AgentGAIA()

    # Example query
    query = (
        "What is the GDP growth difference between France and Germany after "
        "their 2020 policy changes?"
    )

    async def main():
        answer = await agent(query)
        print(answer)

    asyncio.run(main())