"""LangGraph integration for the PAM deploy Agent."""

from __future__ import annotations

from typing import Any, Literal

from .agent import PamDeployAgent

GraphFlow = Literal["global", "deploy"]


def build_langgraph(agent: PamDeployAgent | None = None, flow: GraphFlow = "deploy"):
    try:
        from langgraph.graph import END, START, StateGraph
    except ImportError as exc:  # pragma: no cover - depends on optional package
        raise RuntimeError(
            "langgraph is not installed. Install project dependencies with "
            "`pip install -e .`."
        ) from exc

    runtime = agent or PamDeployAgent()

    def create_state_node(state: dict[str, Any]) -> dict[str, Any]:
        agent_state = runtime.create_state(
            params=state["params"],
            execution_strategy=state.get("execution_strategy", "hybrid_node_mcp"),
            run_id=state.get("run_id"),
            script_entry=state.get("script_entry"),
            config_path=state.get("config_path"),
            trace_file_path=state.get("trace_file_path"),
            target_ips=state.get("target_ips"),
        )
        return {"agent_state": agent_state}

    def run_global_node(state: dict[str, Any]) -> dict[str, Any]:
        agent_state = runtime.run_global_flow(state["agent_state"])
        return {"agent_state": agent_state}

    def run_ip_node(state: dict[str, Any]) -> dict[str, Any]:
        agent_state = runtime.run_ip_flow(state["agent_state"])
        return {"agent_state": agent_state}

    def report_node(state: dict[str, Any]) -> dict[str, Any]:
        return {"report": runtime.render_report(state["agent_state"])}

    graph = StateGraph(dict)
    graph.add_node("create_state", create_state_node)
    graph.add_node("run_global", run_global_node)
    graph.add_node("run_ip", run_ip_node)
    graph.add_node("report", report_node)

    graph.add_edge(START, "create_state")
    graph.add_edge("create_state", "run_global")
    if flow == "global":
        graph.add_edge("run_global", END)
    else:
        graph.add_edge("run_global", "run_ip")
        graph.add_edge("run_ip", "report")
        graph.add_edge("report", END)
    return graph.compile()


def build_graph_or_none(agent: PamDeployAgent | None = None, flow: GraphFlow = "deploy"):
    try:
        return build_langgraph(agent=agent, flow=flow)
    except RuntimeError:
        return None