use std::collections::VecDeque;
use std::io::{BufRead, BufReader, Write};
use std::net::{SocketAddr, TcpListener, TcpStream};
use std::path::PathBuf;
use std::process::{Child, ChildStdin, ChildStdout, Command, ExitCode, Stdio};
use std::sync::mpsc::{self, Receiver};
use std::thread;
use std::time::{Duration, Instant};

use distribution::node::DistributedNodeConfig;
use iroh::EndpointAddr;
use iroh_driver::{IrohDriver, IrohDriverConfig};
use mvp_system::actors::node_agent::{
    NodeAgentActor, NodeAgentMsg, NodeAgentReport, StageCommandWire, StageProvisionWire,
};
use mvp_system::actors::orchestrator::{
    LifecycleEventWire, OrchestratorActor, OrchestratorMsg, OrchestratorReport, RunCommandWire,
    StageRefWire,
};
use mvp_system::actors::register_mvp_actor_codecs;
use mvp_system::dashboard::MvpDashboard;
use mvp_system::distribution_stack::DistributionRuntimeStack;
use mvp_system::driver_pumps;
use mvp_system::observability_surface as obs;
use mvp_system::orchestrator_run_fsm as fsm;
use mvp_system::run_plan as plan;
use mvp_system::stage_controller as stage;
use serde::{Deserialize, Serialize};
use serde_json::json;
use swactor::actor::ActorAddress;
use swactor::runtime::ExternalSender;

const RUN_ID: u64 = 77;
const ORCHESTRATOR_LOGICAL_NODE_ID: u64 = 900;
const NODE0_LOGICAL_ID: u64 = 11;
const NODE1_LOGICAL_ID: u64 = 12;
const MAX_TOKENS: u64 = 1;

fn main() -> ExitCode {
    let args = std::env::args().collect::<Vec<_>>();
    let result = if args.iter().any(|arg| arg == "--role=node") {
        run_node_role(&args)
    } else if std::env::var_os("MVP_DASHBOARD").is_some() {
        run_supervisor_dashboard_loop()
    } else {
        run_supervisor_once(RUN_ID, None, true)
    };

    match result {
        Ok(()) => ExitCode::SUCCESS,
        Err(error) => {
            eprintln!("mvp-local-e2e: {error}");
            ExitCode::from(1)
        }
    }
}

#[derive(Clone, Debug, Serialize, Deserialize)]
struct EdgeFrame {
    edge_id: u64,
    object_id: u64,
    sequence: u64,
    kind: String,
    token_id: Option<u32>,
    eos: bool,
}

#[derive(Clone, Debug, Deserialize)]
struct NodeReady {
    #[serde(rename = "type")]
    kind: String,
    endpoint: EndpointAddr,
    node_actor: ActorAddress,
    token_in_addr: SocketAddr,
    logical_node_id: u64,
    stage_index: u32,
}

#[derive(Clone, Debug, Deserialize)]
struct NodeStdoutLine {
    #[serde(rename = "type")]
    kind: String,
    stage_index: Option<u32>,
    event: Option<String>,
}

fn record_dashboard_event(dashboard: &mut Option<&mut MvpDashboard>, event: obs::Event) {
    if let Some(dashboard) = dashboard.as_deref_mut() {
        dashboard.record_event(event);
    }
}

fn run_event(run_id: u64, kind: obs::EventKind) -> obs::Event {
    obs::Event::RunScoped {
        kind,
        run_id: obs::RunId(run_id),
        reason: None,
        component: obs::Component::Orchestrator,
    }
}

fn run_fault_event(run_id: u64) -> obs::Event {
    obs::Event::RunScoped {
        kind: obs::EventKind::RunFaulted,
        run_id: obs::RunId(run_id),
        reason: Some(obs::FaultReason::WorkerCrashed),
        component: obs::Component::Orchestrator,
    }
}

fn node_event(node_id: u64, kind: obs::EventKind) -> obs::Event {
    obs::Event::NodeScoped {
        kind,
        node_id: obs::NodeId(node_id),
        component: obs::Component::NodeBoot,
    }
}

fn stage_event(run_id: u64, stage_index: u32, kind: obs::EventKind) -> obs::Event {
    obs::Event::StageScoped {
        kind,
        run_id: obs::RunId(run_id),
        stage_index: obs::StageIndex(stage_index),
        reason: if kind == obs::EventKind::StageFaulted {
            Some(obs::FaultReason::WorkerCrashed)
        } else {
            None
        },
        component: obs::Component::StageController,
    }
}

fn object_event(run_id: u64, object_id: u64, sequence: u64, kind: obs::EventKind) -> obs::Event {
    let _ = run_id;
    obs::Event::ObjectScoped {
        kind,
        object_id: obs::ObjectId(object_id),
        sequence: obs::Sequence(sequence),
        component: obs::Component::TokenEndpoint,
    }
}

fn run_supervisor_dashboard_loop() -> Result<(), String> {
    let mut dashboard = MvpDashboard::start_from_env()?;
    eprintln!("mvp-local-e2e: dashboard {}", dashboard.url());
    eprintln!("mvp-local-e2e: MVP_DASHBOARD=1, repeating local scenario until Ctrl+C");
    let mut run_id = RUN_ID;
    loop {
        match run_supervisor_once(run_id, Some(&mut dashboard), false) {
            Ok(()) => eprintln!("mvp-local-e2e: run {run_id} ok"),
            Err(error) => eprintln!("mvp-local-e2e: run {run_id} failed: {error}"),
        }
        run_id = run_id.saturating_add(1);
        thread::sleep(Duration::from_secs(1));
    }
}

fn run_supervisor_once(
    run_id: u64,
    mut dashboard: Option<&mut MvpDashboard>,
    print_summary: bool,
) -> Result<(), String> {
    let _tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?;
    let mut driver = new_driver(_tokio.handle().clone())?;
    let stack = DistributionRuntimeStack::new_with_codecs(
        driver.node_id(),
        DistributedNodeConfig::default(),
        register_mvp_actor_codecs,
    );
    driver.enable_actor_bridge(
        stack.runtime.clone(),
        stack.codec.clone(),
        stack.actor_bridge_routes(),
        stack.actors.swim,
        stack.relay_mirror.clone(),
        stack.route_view.clone(),
    );

    let orchestrator_report = stack
        .runtime
        .new_inbox::<OrchestratorReport>()
        .map_err(|e| format!("orchestrator report inbox: {e}"))?;
    let orchestrator_addr = stack
        .runtime
        .spawn(OrchestratorActor::new(
            fsm::RunConfig {
                run_id: fsm::RunId(run_id),
                max_tokens: MAX_TOKENS,
                prompt: vec![1, 2, 3],
            },
            Some(*orchestrator_report.addr()),
        ))
        .map_err(|e| format!("spawn orchestrator actor: {e}"))?;
    stack.register_local_actor(driver.register_actor(orchestrator_addr, 1));
    record_dashboard_event(
        &mut dashboard,
        node_event(ORCHESTRATOR_LOGICAL_NODE_ID, obs::EventKind::NodeStarted),
    );

    let token_out_listener = TcpListener::bind("127.0.0.1:0")
        .map_err(|e| format!("bind orchestrator token-out listener: {e}"))?;
    let token_out_addr = token_out_listener
        .local_addr()
        .map_err(|e| format!("token-out local addr: {e}"))?;
    spawn_token_out_reader(
        token_out_listener,
        stack.runtime.create_sender(),
        orchestrator_addr,
    );

    let self_endpoint_json = serde_json::to_string(&driver.endpoint_addr())
        .map_err(|e| format!("serialize endpoint addr: {e}"))?;
    let orchestrator_actor_json = serde_json::to_string(&orchestrator_addr)
        .map_err(|e| format!("serialize orchestrator actor: {e}"))?;

    let mut node1 = spawn_node_process(
        NODE1_LOGICAL_ID,
        1,
        7001,
        token_out_addr,
        &self_endpoint_json,
        &orchestrator_actor_json,
    )?;
    let mut node0 = spawn_node_process(
        NODE0_LOGICAL_ID,
        0,
        7000,
        node1.ready.token_in_addr,
        &self_endpoint_json,
        &orchestrator_actor_json,
    )?;
    record_dashboard_event(
        &mut dashboard,
        node_event(NODE1_LOGICAL_ID, obs::EventKind::NodeStarted),
    );
    record_dashboard_event(
        &mut dashboard,
        node_event(NODE1_LOGICAL_ID, obs::EventKind::NodeAvailable),
    );
    record_dashboard_event(
        &mut dashboard,
        node_event(NODE0_LOGICAL_ID, obs::EventKind::NodeStarted),
    );
    record_dashboard_event(
        &mut dashboard,
        node_event(NODE0_LOGICAL_ID, obs::EventKind::NodeAvailable),
    );

    let started_at = Instant::now();
    wait_for_routes(
        &mut driver,
        &stack,
        &[node0.ready.node_actor, node1.ready.node_actor],
        Duration::from_secs(20),
    )?;

    let run_plan = two_stage_plan(run_id)?;
    stack
        .runtime
        .send_to(
            orchestrator_addr,
            OrchestratorMsg::ObservePoolReady {
                nodes: vec![NODE0_LOGICAL_ID, NODE1_LOGICAL_ID],
            },
        )
        .map_err(|e| format!("observe pool ready: {e}"))?;
    record_dashboard_event(&mut dashboard, run_event(run_id, obs::EventKind::PoolReady));
    stack
        .runtime
        .send_to(
            orchestrator_addr,
            OrchestratorMsg::ObservePlanAvailable {
                run_id,
                stages: vec![
                    StageRefWire {
                        stage_index: 0,
                        node_id: NODE0_LOGICAL_ID,
                    },
                    StageRefWire {
                        stage_index: 1,
                        node_id: NODE1_LOGICAL_ID,
                    },
                ],
            },
        )
        .map_err(|e| format!("observe plan: {e}"))?;
    record_dashboard_event(
        &mut dashboard,
        run_event(run_id, obs::EventKind::RunPlanned),
    );
    stack
        .runtime
        .send_to(
            orchestrator_addr,
            OrchestratorMsg::ObserveTokenInEndpointReady,
        )
        .map_err(|e| format!("observe token-in endpoint: {e}"))?;
    stack
        .runtime
        .send_to(
            orchestrator_addr,
            OrchestratorMsg::ObserveTokenOutEndpointReady,
        )
        .map_err(|e| format!("observe token-out endpoint: {e}"))?;
    record_dashboard_event(
        &mut dashboard,
        run_event(run_id, obs::EventKind::ReadinessBarrierPassed),
    );

    let mut injected = false;
    let mut completed = false;
    let mut torn_down = false;
    let mut stage_ready_count = 0usize;
    let mut token_received = false;
    let mut sent_stop_to_node0 = false;
    let mut sent_stop_to_node1 = false;

    while started_at.elapsed() < Duration::from_secs(30) {
        pump_network(&mut driver, &stack);
        stage_ready_count += drain_node_stdout(&node0.stdout_rx, run_id, &mut dashboard);
        stage_ready_count += drain_node_stdout(&node1.stdout_rx, run_id, &mut dashboard);

        while let Some(report) = orchestrator_report.try_recv() {
            match report {
                OrchestratorReport::Command(command) => match command {
                    RunCommandWire::ProvisionStage { stage_index, .. } => {
                        let provision = stage_provision_wire(&run_plan, stage_index)?;
                        let target = if stage_index == 0 {
                            node0.ready.node_actor
                        } else {
                            node1.ready.node_actor
                        };
                        stack
                            .runtime
                            .send_to(target, NodeAgentMsg::ProvisionStage(provision))
                            .map_err(|e| format!("send provision to stage {stage_index}: {e}"))?;
                        record_dashboard_event(
                            &mut dashboard,
                            stage_event(run_id, stage_index, obs::EventKind::StageProvisionStarted),
                        );
                    }
                    RunCommandWire::InjectPrompt {
                        sequence, prompt, ..
                    } => {
                        injected = true;
                        record_dashboard_event(
                            &mut dashboard,
                            object_event(
                                run_id,
                                9_000 + sequence,
                                sequence,
                                obs::EventKind::PromptInjected,
                            ),
                        );
                        write_edge_frame(
                            node0.ready.token_in_addr,
                            EdgeFrame {
                                edge_id: 7000,
                                object_id: 9_000 + sequence,
                                sequence,
                                kind: "token".to_owned(),
                                token_id: prompt.first().copied(),
                                eos: false,
                            },
                        )?;
                    }
                    RunCommandWire::StopRun {
                        stage_index,
                        run_id,
                    } => {
                        let target = if stage_index == 0 {
                            sent_stop_to_node0 = true;
                            node0.ready.node_actor
                        } else {
                            sent_stop_to_node1 = true;
                            node1.ready.node_actor
                        };
                        stack
                            .runtime
                            .send_to(target, NodeAgentMsg::StopRun { run_id })
                            .map_err(|e| format!("send stop to stage {stage_index}: {e}"))?;
                        record_dashboard_event(
                            &mut dashboard,
                            stage_event(run_id, stage_index, obs::EventKind::StopRunSent),
                        );
                    }
                    RunCommandWire::TearDownTokenEndpoints { .. } => {
                        stack
                            .runtime
                            .send_to(
                                orchestrator_addr,
                                OrchestratorMsg::ObserveTokenEndpointsStopped,
                            )
                            .map_err(|e| format!("observe token endpoints stopped: {e}"))?;
                    }
                    RunCommandWire::CreateTokenInEndpoint { .. }
                    | RunCommandWire::CreateTokenOutEndpoint { .. }
                    | RunCommandWire::BroadcastStart { .. } => {}
                },
                OrchestratorReport::Lifecycle(event) => match event {
                    LifecycleEventWire::RunCompleted { run_id } => {
                        token_received = true;
                        record_dashboard_event(
                            &mut dashboard,
                            object_event(run_id, 9_000, 0, obs::EventKind::TokenReceived),
                        );
                        record_dashboard_event(
                            &mut dashboard,
                            run_event(run_id, obs::EventKind::RunCompleted),
                        );
                        completed = true;
                    }
                    LifecycleEventWire::RunTornDown { run_id } => {
                        record_dashboard_event(
                            &mut dashboard,
                            run_event(run_id, obs::EventKind::RunTornDown),
                        );
                        torn_down = true;
                    }
                    LifecycleEventWire::RunRejected { run_id }
                    | LifecycleEventWire::RunFaulted { run_id } => {
                        record_dashboard_event(&mut dashboard, run_fault_event(run_id));
                        return Err(format!("run failed: {event:?}"));
                    }
                },
                OrchestratorReport::Snapshot { .. } => {}
            }
        }

        token_received |= completed;

        if injected && completed && torn_down && sent_stop_to_node0 && sent_stop_to_node1 {
            shutdown_node(&mut node0);
            shutdown_node(&mut node1);
            let summary = json!({
                "ok": true,
                "actor_plane": "iroh-swactor",
                "processes": {
                    "orchestrator": std::process::id(),
                    "node0": node0.child.id(),
                    "node1": node1.child.id(),
                },
                "node0_endpoint": node0.ready.endpoint,
                "node1_endpoint": node1.ready.endpoint,
                "node0_logical_id": node0.ready.logical_node_id,
                "node1_logical_id": node1.ready.logical_node_id,
                "node0_stage_index": node0.ready.stage_index,
                "node1_stage_index": node1.ready.stage_index,
                "data_plane": "tcp-loopback-streams",
                "worker_processes": "mvp-dumb-worker-per-node",
                "injected_prompt_observed": injected,
                "token_received_observed": token_received,
                "run_completed_observed": completed,
                "run_torn_down_observed": torn_down,
                "stop_sent_to_all_nodes": sent_stop_to_node0 && sent_stop_to_node1,
                "node_route_count": stack.route_view.read().map(|view| view.len()).unwrap_or_default(),
                "stage_ready_stdout_count": stage_ready_count,
            });
            if print_summary {
                println!("{summary}");
            } else {
                eprintln!("mvp-local-e2e: run {run_id} summary {summary}");
            }
            return Ok(());
        }

        thread::sleep(Duration::from_millis(10));
    }

    record_dashboard_event(&mut dashboard, run_fault_event(run_id));
    shutdown_node(&mut node0);
    shutdown_node(&mut node1);
    Err(format!(
        "timed out: injected={injected} completed={completed} torn_down={torn_down} stop0={sent_stop_to_node0} stop1={sent_stop_to_node1}"
    ))
}

fn run_node_role(args: &[String]) -> Result<(), String> {
    let logical_node_id = parse_arg(args, "--logical-node-id")?
        .parse::<u64>()
        .map_err(|e| format!("logical node id: {e}"))?;
    let stage_index = parse_arg(args, "--stage-index")?
        .parse::<u32>()
        .map_err(|e| format!("stage index: {e}"))?;
    let inbound_edge_id = parse_arg(args, "--inbound-edge-id")?
        .parse::<u64>()
        .map_err(|e| format!("inbound edge id: {e}"))?;
    let outbound_addr = parse_arg(args, "--outbound-addr")?
        .parse::<SocketAddr>()
        .map_err(|e| format!("outbound addr: {e}"))?;
    let seed: EndpointAddr = serde_json::from_str(parse_arg(args, "--seed-endpoint")?)
        .map_err(|e| format!("seed endpoint json: {e}"))?;
    let orchestrator_addr: ActorAddress =
        serde_json::from_str(parse_arg(args, "--orchestrator-actor")?)
            .map_err(|e| format!("orchestrator actor json: {e}"))?;

    let _tokio = tokio::runtime::Runtime::new().map_err(|e| format!("tokio runtime: {e}"))?;
    let mut driver = new_driver(_tokio.handle().clone())?;
    let stack = DistributionRuntimeStack::new_with_codecs(
        driver.node_id(),
        DistributedNodeConfig::default(),
        register_mvp_actor_codecs,
    );
    driver.enable_actor_bridge(
        stack.runtime.clone(),
        stack.codec.clone(),
        stack.actor_bridge_routes(),
        stack.actors.swim,
        stack.relay_mirror.clone(),
        stack.route_view.clone(),
    );
    driver.join(&[seed]);

    let node_report = stack
        .runtime
        .new_inbox::<NodeAgentReport>()
        .map_err(|e| format!("node report inbox: {e}"))?;
    let node_actor = stack
        .runtime
        .spawn(NodeAgentActor::new(
            stage::NodeId(logical_node_id),
            orchestrator_addr,
            Some(*node_report.addr()),
        ))
        .map_err(|e| format!("spawn node actor: {e}"))?;
    stack.register_local_actor(driver.register_actor(node_actor, 1));

    let listener = TcpListener::bind("127.0.0.1:0")
        .map_err(|e| format!("bind node inbound edge listener: {e}"))?;
    let token_in_addr = listener
        .local_addr()
        .map_err(|e| format!("node inbound local addr: {e}"))?;
    spawn_inbound_edge_reader(
        listener,
        inbound_edge_id,
        stack.runtime.create_sender(),
        node_actor,
    );

    println!(
        "{}",
        json!({
            "type": "ready",
            "role": "node",
            "endpoint": driver.endpoint_addr(),
            "node_actor": node_actor,
            "token_in_addr": token_in_addr,
            "logical_node_id": logical_node_id,
            "stage_index": stage_index,
        })
    );
    std::io::stdout()
        .flush()
        .map_err(|e| format!("flush ready: {e}"))?;

    let (stdin_tx, stdin_rx) = mpsc::channel::<serde_json::Value>();
    thread::spawn(move || {
        let stdin = std::io::stdin();
        for line in stdin.lock().lines().map_while(Result::ok) {
            if let Ok(value) = serde_json::from_str::<serde_json::Value>(&line) {
                let _ = stdin_tx.send(value);
            }
        }
    });

    let mut worker = WorkerProc::spawn()?;
    let mut pending_commands = VecDeque::new();
    let mut outbound_stream: Option<TcpStream> = None;
    let started_at = Instant::now();

    loop {
        if let Ok(value) = stdin_rx.try_recv() {
            if value.get("type").and_then(|value| value.as_str()) == Some("shutdown") {
                let _ = worker.shutdown();
                return Ok(());
            }
        }

        pump_network(&mut driver, &stack);

        while let Some(report) = node_report.try_recv() {
            match report {
                NodeAgentReport::Command(command) => pending_commands.push_back(command),
                NodeAgentReport::Lifecycle(event) => {
                    println!(
                        "{}",
                        json!({
                            "type": "node_lifecycle",
                            "stage_index": stage_index,
                            "event": format!("{event:?}"),
                        })
                    );
                    let _ = std::io::stdout().flush();
                }
                NodeAgentReport::Snapshot { .. } => {}
            }
        }

        let route_to_orchestrator_ready = stack
            .route_view
            .read()
            .map(|view| view.contains_key(&orchestrator_addr))
            .unwrap_or(false);
        let mut deferred = VecDeque::new();
        while let Some(command) = pending_commands.pop_front() {
            if !route_to_orchestrator_ready && readiness_command(&command) {
                deferred.push_back(command);
                continue;
            }
            handle_node_command(
                command,
                &stack.runtime,
                node_actor,
                &mut worker,
                outbound_addr,
                &mut outbound_stream,
            )?;
        }
        pending_commands = deferred;

        if started_at.elapsed() > Duration::from_secs(60) {
            return Err("node role timed out".to_owned());
        }
        thread::sleep(Duration::from_millis(10));
    }
}

fn new_driver(handle: tokio::runtime::Handle) -> Result<IrohDriver, String> {
    IrohDriver::with_handle(
        handle,
        IrohDriverConfig {
            secret_key: None,
            relay_mode: iroh::RelayMode::Disabled,
            node: DistributedNodeConfig::default(),
            peer_auth: None,
            additional_alpns: vec![],
        },
    )
    .map_err(|e| format!("create iroh driver: {e}"))
}

fn pump_network(driver: &mut IrohDriver, stack: &DistributionRuntimeStack) {
    stack.tick_protocol_actors(Instant::now());
    driver.pump_inbound_to_actors();
    stack.pump_runtime_once();
    driver.drain_outbox(&stack.outbox);
}

fn wait_for_routes(
    driver: &mut IrohDriver,
    stack: &DistributionRuntimeStack,
    actors: &[ActorAddress],
    timeout: Duration,
) -> Result<(), String> {
    let started = Instant::now();
    while started.elapsed() < timeout {
        pump_network(driver, stack);
        let ready = stack
            .route_view
            .read()
            .map(|view| actors.iter().all(|actor| view.contains_key(actor)))
            .unwrap_or(false);
        if ready {
            return Ok(());
        }
        thread::sleep(Duration::from_millis(20));
    }
    Err("directory routes for node actors did not converge".to_owned())
}

fn two_stage_plan(run_id: u64) -> Result<plan::RunPlan, String> {
    plan::plan_run(plan::PlannerInput {
        run_id: plan::RunId(run_id),
        orchestrator_node_id: plan::NodeId(ORCHESTRATOR_LOGICAL_NODE_ID),
        model: plan::ModelFacts {
            model_id: "local-e2e-fixture".to_owned(),
            num_layers: 4,
            hidden_dim: 8,
            dtype_family: plan::DTypeFamily::BFloat,
            dtype_width_bytes: 2,
            max_seq_len: 8,
            eos_token_id: 99,
        },
        runtime: plan::RuntimeConfig { max_tokens: 1 },
        candidate_pool: vec![
            plan::NodeId(NODE0_LOGICAL_ID),
            plan::NodeId(NODE1_LOGICAL_ID),
        ],
        stage_count: 2,
        placement: plan::PlacementInput::FixedLinear(vec![
            plan::StagePlacement {
                stage_index: 0,
                node_id: plan::NodeId(NODE0_LOGICAL_ID),
            },
            plan::StagePlacement {
                stage_index: 1,
                node_id: plan::NodeId(NODE1_LOGICAL_ID),
            },
        ]),
        activation_ring: plan::RingSpec::test_default_activation(),
        token_ring: plan::RingSpec::test_default_token(),
    })
    .map_err(|e| format!("plan rejected: {:?}", e.kind()))
}

fn stage_provision_wire(
    plan: &plan::RunPlan,
    stage_index: u32,
) -> Result<StageProvisionWire, String> {
    let provision = plan::derive_stage_provision(plan, stage_index)
        .map_err(|e| format!("derive stage provision {stage_index}: {e:?}"))?;
    Ok(StageProvisionWire {
        run_id: provision.run_id.0,
        authorized_orchestrator: ORCHESTRATOR_LOGICAL_NODE_ID,
        node_id: provision.node_id.0,
        stage_index: provision.stage_index,
        stage_count: provision.stage_count,
        layer_start: provision.layer_start,
        layer_end_exclusive: provision.layer_end_exclusive,
        inbound_edge_id: provision.inbound.edge_id.0,
        outbound_edge_id: provision.outbound.edge_id.0,
        weight_artifact: "local-e2e-fixture".to_owned(),
    })
}

fn write_edge_frame(addr: SocketAddr, frame: EdgeFrame) -> Result<(), String> {
    let mut stream = TcpStream::connect(addr).map_err(|e| format!("connect edge {addr}: {e}"))?;
    stream
        .write_all(&driver_pumps::encode_edge_preamble(driver_pumps::EdgeId(
            frame.edge_id,
        )))
        .map_err(|e| format!("write edge preamble: {e}"))?;
    write_json_frame(&mut stream, &frame)
}

fn write_json_frame(stream: &mut TcpStream, frame: &EdgeFrame) -> Result<(), String> {
    serde_json::to_writer(&mut *stream, frame).map_err(|e| format!("serialize edge frame: {e}"))?;
    stream
        .write_all(b"\n")
        .map_err(|e| format!("write edge frame newline: {e}"))?;
    stream.flush().map_err(|e| format!("flush edge frame: {e}"))
}

fn spawn_inbound_edge_reader(
    listener: TcpListener,
    expected_edge_id: u64,
    sender: ExternalSender,
    node_actor: ActorAddress,
) {
    thread::spawn(move || {
        for incoming in listener.incoming() {
            let Ok(mut stream) = incoming else { continue };
            let mut preamble = [0u8; 8];
            if std::io::Read::read_exact(&mut stream, &mut preamble).is_err() {
                continue;
            }
            if preamble != expected_edge_id.to_le_bytes() {
                continue;
            }
            let mut reader = BufReader::new(stream);
            loop {
                let mut line = String::new();
                match reader.read_line(&mut line) {
                    Ok(0) => break,
                    Ok(_) => {
                        if let Ok(frame) = serde_json::from_str::<EdgeFrame>(&line) {
                            let _ = sender.send_to(
                                node_actor,
                                NodeAgentMsg::ObjectLoaded {
                                    edge_id: frame.edge_id,
                                    object_id: frame.object_id,
                                    sequence: frame.sequence,
                                    handle_generation: 1,
                                    handle_id: frame.object_id,
                                },
                            );
                        }
                    }
                    Err(_) => break,
                }
            }
        }
    });
}

fn spawn_token_out_reader(
    listener: TcpListener,
    sender: ExternalSender,
    orchestrator_addr: ActorAddress,
) {
    thread::spawn(move || {
        for incoming in listener.incoming() {
            let Ok(mut stream) = incoming else { continue };
            let mut preamble = [0u8; 8];
            if std::io::Read::read_exact(&mut stream, &mut preamble).is_err() {
                continue;
            }
            let mut reader = BufReader::new(stream);
            let mut line = String::new();
            if reader.read_line(&mut line).is_ok()
                && let Ok(frame) = serde_json::from_str::<EdgeFrame>(&line)
            {
                let _ = sender.send_to(
                    orchestrator_addr,
                    OrchestratorMsg::ObserveTokenReceived {
                        sequence: frame.sequence,
                        token_id: frame.token_id.unwrap_or(99),
                        eos: frame.eos,
                    },
                );
            }
        }
    });
}

fn handle_node_command(
    command: StageCommandWire,
    runtime: &swactor::runtime::Runtime,
    node_actor: ActorAddress,
    worker: &mut WorkerProc,
    outbound_addr: SocketAddr,
    outbound_stream: &mut Option<TcpStream>,
) -> Result<(), String> {
    match command {
        StageCommandWire::EstablishInboundEdge { edge_id } => runtime
            .send_to(node_actor, NodeAgentMsg::MarkInboundEdgeReady { edge_id })
            .map_err(|e| format!("mark inbound ready: {e}")),
        StageCommandWire::EstablishOutboundEdge { edge_id } => {
            let mut stream = TcpStream::connect(outbound_addr)
                .map_err(|e| format!("connect outbound edge {edge_id} to {outbound_addr}: {e}"))?;
            stream
                .write_all(&driver_pumps::encode_edge_preamble(driver_pumps::EdgeId(
                    edge_id,
                )))
                .map_err(|e| format!("write outbound preamble: {e}"))?;
            *outbound_stream = Some(stream);
            runtime
                .send_to(node_actor, NodeAgentMsg::MarkOutboundEdgeReady { edge_id })
                .map_err(|e| format!("mark outbound ready: {e}"))
        }
        StageCommandWire::ConfigureWorkerRole { .. } => {
            worker.initialize()?;
            runtime
                .send_to(node_actor, NodeAgentMsg::MarkWorkerReady)
                .map_err(|e| format!("mark worker ready: {e}"))
        }
        StageCommandWire::LoadWeights { .. } => runtime
            .send_to(node_actor, NodeAgentMsg::MarkWeightsReady)
            .map_err(|e| format!("mark weights ready: {e}")),
        StageCommandWire::ExecuteStep {
            step_id,
            sequence,
            output_edge_ids,
            ..
        } => {
            worker.execute_step(step_id)?;
            runtime
                .send_to(node_actor, NodeAgentMsg::StepCompleted { step_id })
                .map_err(|e| format!("mark step completed: {e}"))?;
            let output_edge_id = output_edge_ids.first().copied().unwrap_or(0);
            let stream = outbound_stream
                .as_mut()
                .ok_or_else(|| "outbound stream missing for ExecuteStep".to_owned())?;
            write_json_frame(
                stream,
                &EdgeFrame {
                    edge_id: output_edge_id,
                    object_id: 10_000 + step_id,
                    sequence,
                    kind: if output_edge_id == 7002 {
                        "token".to_owned()
                    } else {
                        "activation".to_owned()
                    },
                    token_id: (output_edge_id == 7002).then_some(99),
                    eos: output_edge_id == 7002,
                },
            )
        }
        StageCommandWire::ReleaseInputHandle { .. }
        | StageCommandWire::StopLocalEdges { .. }
        | StageCommandWire::ReleaseRunDeviceObjects { .. }
        | StageCommandWire::RewireEdge { .. } => Ok(()),
    }
}

fn readiness_command(command: &StageCommandWire) -> bool {
    matches!(
        command,
        StageCommandWire::EstablishInboundEdge { .. }
            | StageCommandWire::EstablishOutboundEdge { .. }
            | StageCommandWire::ConfigureWorkerRole { .. }
            | StageCommandWire::LoadWeights { .. }
    )
}

struct WorkerProc {
    child: Child,
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
}

impl WorkerProc {
    fn spawn() -> Result<Self, String> {
        let mut child = Command::new(worker_bin_path())
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .stderr(Stdio::null())
            .spawn()
            .map_err(|e| format!("spawn dumb worker: {e}"))?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| "worker stdin missing".to_owned())?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| "worker stdout missing".to_owned())?;
        Ok(Self {
            child,
            stdin,
            stdout: BufReader::new(stdout),
        })
    }

    fn initialize(&mut self) -> Result<(), String> {
        self.command(
            json!({"type":"InitializeWorker","helper_abi_version":1}),
            "WorkerReady",
        )
    }

    fn execute_step(&mut self, step_id: u64) -> Result<(), String> {
        self.command(
            json!({"type":"ExecuteStep","step_id":step_id}),
            "StepCompleted",
        )
    }

    fn shutdown(&mut self) -> Result<(), String> {
        self.command(json!({"type":"ShutdownWorker"}), "WorkerStopped")
    }

    fn command(&mut self, command: serde_json::Value, expected: &str) -> Result<(), String> {
        writeln!(self.stdin, "{command}").map_err(|e| format!("write worker command: {e}"))?;
        self.stdin
            .flush()
            .map_err(|e| format!("flush worker stdin: {e}"))?;
        let mut line = String::new();
        self.stdout
            .read_line(&mut line)
            .map_err(|e| format!("read worker stdout: {e}"))?;
        let value: serde_json::Value = serde_json::from_str(&line)
            .map_err(|e| format!("parse worker stdout {line:?}: {e}"))?;
        if value.get("type").and_then(|value| value.as_str()) == Some(expected) {
            Ok(())
        } else {
            Err(format!("worker emitted {value}, expected {expected}"))
        }
    }
}

impl Drop for WorkerProc {
    fn drop(&mut self) {
        let _ = self.child.try_wait();
    }
}

struct NodeChild {
    child: Child,
    stdin: ChildStdin,
    stdout_rx: Receiver<NodeStdoutLine>,
    ready: NodeReady,
}

fn spawn_node_process(
    logical_node_id: u64,
    stage_index: u32,
    inbound_edge_id: u64,
    outbound_addr: SocketAddr,
    seed_endpoint_json: &str,
    orchestrator_actor_json: &str,
) -> Result<NodeChild, String> {
    let mut child = Command::new(std::env::current_exe().map_err(|e| format!("current exe: {e}"))?)
        .arg("--role=node")
        .arg("--logical-node-id")
        .arg(logical_node_id.to_string())
        .arg("--stage-index")
        .arg(stage_index.to_string())
        .arg("--inbound-edge-id")
        .arg(inbound_edge_id.to_string())
        .arg("--outbound-addr")
        .arg(outbound_addr.to_string())
        .arg("--seed-endpoint")
        .arg(seed_endpoint_json)
        .arg("--orchestrator-actor")
        .arg(orchestrator_actor_json)
        .stdin(Stdio::piped())
        .stdout(Stdio::piped())
        .stderr(Stdio::piped())
        .spawn()
        .map_err(|e| format!("spawn node {stage_index}: {e}"))?;

    let stdin = child
        .stdin
        .take()
        .ok_or_else(|| "node stdin missing".to_owned())?;
    let stdout = child
        .stdout
        .take()
        .ok_or_else(|| "node stdout missing".to_owned())?;
    let mut reader = BufReader::new(stdout);
    let mut ready_line = String::new();
    reader
        .read_line(&mut ready_line)
        .map_err(|e| format!("read node ready: {e}"))?;
    let ready: NodeReady = serde_json::from_str(&ready_line)
        .map_err(|e| format!("parse node ready {ready_line:?}: {e}"))?;
    if ready.kind != "ready" {
        return Err(format!("node first line was not ready: {ready_line}"));
    }

    let (tx, rx) = mpsc::channel();
    thread::spawn(move || {
        for line in reader.lines().map_while(Result::ok) {
            if let Ok(value) = serde_json::from_str::<NodeStdoutLine>(&line) {
                let _ = tx.send(value);
            }
        }
    });

    Ok(NodeChild {
        child,
        stdin,
        stdout_rx: rx,
        ready,
    })
}

fn shutdown_node(node: &mut NodeChild) {
    let _ = writeln!(node.stdin, "{}", json!({"type":"shutdown"}));
    let _ = node.stdin.flush();
    let started = Instant::now();
    while started.elapsed() < Duration::from_secs(3) {
        if matches!(node.child.try_wait(), Ok(Some(_))) {
            return;
        }
        thread::sleep(Duration::from_millis(20));
    }
    let _ = node.child.kill();
    let _ = node.child.wait();
}

fn drain_node_stdout(
    rx: &Receiver<NodeStdoutLine>,
    run_id: u64,
    dashboard: &mut Option<&mut MvpDashboard>,
) -> usize {
    let mut stage_ready_count = 0;
    while let Ok(line) = rx.try_recv() {
        if line.kind != "node_lifecycle" {
            continue;
        }
        let Some(stage_index) = line.stage_index else {
            continue;
        };
        let Some(event) = line.event.as_deref() else {
            continue;
        };
        if event.contains("StageReady") {
            stage_ready_count += 1;
            record_dashboard_event(
                dashboard,
                stage_event(run_id, stage_index, obs::EventKind::StageReady),
            );
        } else if event.contains("StageStopped") {
            record_dashboard_event(
                dashboard,
                stage_event(run_id, stage_index, obs::EventKind::StageStopped),
            );
        } else if event.contains("StageFault") {
            record_dashboard_event(
                dashboard,
                stage_event(run_id, stage_index, obs::EventKind::StageFaulted),
            );
        }
    }
    stage_ready_count
}

fn parse_arg<'a>(args: &'a [String], name: &str) -> Result<&'a str, String> {
    let index = args
        .iter()
        .position(|arg| arg == name)
        .ok_or_else(|| format!("missing {name}"))?;
    args.get(index + 1)
        .map(String::as_str)
        .ok_or_else(|| format!("missing value for {name}"))
}

fn worker_bin_path() -> PathBuf {
    if let Ok(path) = std::env::var("MVP_DUMB_WORKER_BIN") {
        return PathBuf::from(path);
    }
    let mut path = std::env::current_exe().expect("current exe");
    path.set_file_name(format!("mvp-dumb-worker{}", std::env::consts::EXE_SUFFIX));
    path
}