"""
Graph assembly — wires together the supervisor and worker nodes.

This is the LangGraph equivalent of Loki's runtime agent execution engine
(src/supervisor/mod.rs + src/config/request_context.rs).

In Loki, the runtime:
    1. Loads the agent config (config.yaml)
    2. Compiles tools (tools.sh → binary)
    3. Starts a chat loop: user → LLM → tool calls → LLM → ...
    4. For orchestrators with can_spawn_agents: true, the supervisor module
       manages child agent lifecycle (spawn, check, collect, cancel).

In LangGraph, all of this is declarative:
    1. Define nodes (supervisor, explore, oracle, coder)
    2. Define edges (workers always return to supervisor)
    3. Compile the graph (with optional checkpointer for persistence)
    4. Invoke with initial state

The graph topology:

    ┌─────────────────────────────────────────────┐
    │               SUPERVISOR                     │
    │  (classifies intent, routes to workers)       │
    └─────┬──────────┬──────────┬─────────────────┘
          │          │          │
          ▼          ▼          ▼
      ┌────────┐ ┌────────┐ ┌────────┐
      │EXPLORE │ │ ORACLE │ │ CODER  │
      │(search)│ │(advise)│ │(build) │
      └───┬────┘ └───┬────┘ └───┬────┘
          │          │          │
          └──────────┼──────────┘
                     │
              (back to supervisor)

Every worker returns to the supervisor.  The supervisor decides what to do next:
route to another worker, or end the graph.
"""

from __future__ import annotations

from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import END, START, StateGraph

from sisyphus_langchain.agents.coder import create_coder_node
from sisyphus_langchain.agents.explore import create_explore_node
from sisyphus_langchain.agents.oracle import create_oracle_node
from sisyphus_langchain.agents.supervisor import create_supervisor_node
from sisyphus_langchain.state import SisyphusState


def build_graph(
    *,
    supervisor_model: str = "gpt-4o",
    explore_model: str = "gpt-4o-mini",
    oracle_model: str = "gpt-4o",
    coder_model: str = "gpt-4o",
    use_checkpointer: bool = True,
):
    """
    Build and compile the Sisyphus LangGraph.

    This is the main entry point for creating the agent system.  It wires
    together all nodes and edges, optionally adds a checkpointer for
    persistence, and returns a compiled graph ready to invoke.

    Args:
        supervisor_model: Model for the routing supervisor.
        explore_model: Model for the explore agent (can be cheaper).
        oracle_model: Model for the oracle agent (should be strong).
        coder_model: Model for the coder agent.
        use_checkpointer: Whether to add MemorySaver for session persistence.

    Returns:
        A compiled LangGraph ready to .invoke() or .stream().

    Model cost optimization (mirrors Loki's per-agent model config):
        - supervisor: expensive (accurate routing is critical)
        - explore: cheap (just searching, not reasoning deeply)
        - oracle: expensive (deep reasoning, architecture advice)
        - coder: expensive (writing correct code matters)
    """
    # Create the graph builder with our typed state
    builder = StateGraph(SisyphusState)

    # ── Register nodes ─────────────────────────────────────────────────
    # Each node is a function that takes state and returns state updates.
    # This mirrors Loki's agent registration (agents are discovered by
    # their config.yaml in the agents/ directory).
    builder.add_node("supervisor", create_supervisor_node(supervisor_model))
    builder.add_node("explore", create_explore_node(explore_model))
    builder.add_node("oracle", create_oracle_node(oracle_model))
    builder.add_node("coder", create_coder_node(coder_model))

    # ── Define edges ───────────────────────────────────────────────────
    # Entry point: every invocation starts at the supervisor
    builder.add_edge(START, "supervisor")

    # Workers always return to supervisor (the hub-and-spoke pattern).
    # In Loki, this is implicit: agent__collect returns output to the parent,
    # and the parent (sisyphus) decides what to do next.
    builder.add_edge("explore", "supervisor")
    builder.add_edge("oracle", "supervisor")
    builder.add_edge("coder", "supervisor")

    # The supervisor node itself uses Command(goto=...) to route,
    # so we don't need add_conditional_edges — the Command API
    # handles dynamic routing internally.

    # ── Compile ────────────────────────────────────────────────────────
    checkpointer = MemorySaver() if use_checkpointer else None
    graph = builder.compile(checkpointer=checkpointer)

    return graph