zc2 0.0.28

//! Local compute-node *mesh* orchestration for `zc mesh`.
//!
//! Launches a Docker-based mesh of zakuro compute nodes on a shared network,
//! starts the zakuro worker on each, and runs a broker that discovers them as a
//! fleet (`ZAKURO_PEERS`). This is the scripted form of the manual
//! `docker run …` dance — local mode, no auth.
//!
//! Tunables via env: `ZAKURO_MESH_IMAGE` (node image, default
//! `zakuro-image-compute:latest`), `ZAKURO_MESH_BROKER_IMAGE` (default
//! `zakuroai/zc:latest`), `ZAKURO_MESH_PLATFORM` (default `linux/amd64`).

use std::process::Command;
use std::time::Duration;

const NETWORK: &str = "zk-mesh";
const BROKER: &str = "zk-broker";
const DASH_BASE: u16 = 8090; // node i dashboard -> host 8090+i
const BROKER_PORT: u16 = 9000;
const WORKER_PORT: u16 = 3960;

fn node_image() -> String {
    std::env::var("ZAKURO_MESH_IMAGE").unwrap_or_else(|_| "zakuro-image-compute:latest".to_string())
}
fn broker_image() -> String {
    std::env::var("ZAKURO_MESH_BROKER_IMAGE").unwrap_or_else(|_| "zakuroai/zc:latest".to_string())
}
fn platform() -> Option<String> {
    match std::env::var("ZAKURO_MESH_PLATFORM") {
        Ok(p) if p.trim().is_empty() => None,
        Ok(p) => Some(p),
        Err(_) => Some("linux/amd64".to_string()),
    }
}

// Per-node resource caps — keep the mesh footprint small. Tunable via env.
fn cpus() -> String {
    std::env::var("ZAKURO_MESH_CPUS").unwrap_or_else(|_| "2".to_string())
}
fn memory() -> String {
    std::env::var("ZAKURO_MESH_MEM").unwrap_or_else(|_| "2g".to_string())
}

// --- p2p -------------------------------------------------------------------
// Shared secret used to authenticate peer-broker RPC in p2p mode.
fn peer_key() -> String {
    std::env::var("ZAKURO_MESH_PEER_KEY").unwrap_or_else(|_| "zakuro-mesh".to_string())
}
// p2p is enabled explicitly via env. It used to switch on implicitly whenever
// a mesh auth key was exported, which is gone -- nodes now join the WireGuard
// mesh on the host, not by provisioning a tunnel per container.
fn p2p_enabled() -> bool {
    matches!(
        std::env::var("ZAKURO_MESH_P2P").as_deref(),
        Ok("1") | Ok("true")
    )
}

/// Run `docker <args>`, returning trimmed stdout or an error string.
fn docker(args: &[String]) -> Result<String, String> {
    let out = Command::new("docker")
        .args(args)
        .output()
        .map_err(|e| format!("docker not available: {}", e))?;
    if out.status.success() {
        Ok(String::from_utf8_lossy(&out.stdout).trim().to_string())
    } else {
        Err(String::from_utf8_lossy(&out.stderr).trim().to_string())
    }
}

fn d(args: &[&str]) -> Result<String, String> {
    docker(&args.iter().map(|s| s.to_string()).collect::<Vec<_>>())
}

fn node_name(i: usize) -> String {
    format!("zk{}", i)
}

/// Launch (or relaunch) a mesh of `nodes` compute nodes + a broker.
pub fn up(nodes: usize) {
    let nodes = nodes.clamp(1, 16);
    let image = node_image();
    let plat = platform();
    let p2p = p2p_enabled();
    if p2p {
        println!("  zc mesh — P2P mode (ZAKURO_P2P=true, peer-key set)");
    }
    println!(
        "  zc mesh — launching {} node(s) from {} ({} cpu / {} mem each)",
        nodes,
        image,
        cpus(),
        memory()
    );

    // Shared network (idempotent).
    let _ = d(&["network", "create", NETWORK]);

    for i in 0..nodes {
        let name = node_name(i);
        let dash = (DASH_BASE + i as u16).to_string();
        let _ = d(&["rm", "-f", &name]);

        let mut args: Vec<String> = vec!["run".into(), "-d".into(), "--name".into(), name.clone()];
        if let Some(p) = &plat {
            args.push("--platform".into());
            args.push(p.clone());
        }
        // Cap CPU + memory so the mesh stays light.
        args.extend(["--cpus".into(), cpus(), "--memory".into(), memory()]);
        args.extend([
            "--hostname".into(),
            name.clone(),
            "--network".into(),
            NETWORK.into(),
            "--network-alias".into(),
            name.clone(),
            "-p".into(),
            format!("{}:8080", dash),
        ]);
        // The primary node also publishes the cluster UIs.
        if i == 0 {
            for pm in [
                "8265:8265",
                "8787:8787",
                "7077:7077",
                "8082:8082",
                "8888:8888",
            ] {
                args.push("-p".into());
                args.push(pm.into());
            }
        }
        args.extend([
            "-e".into(),
            "RAY_DASHBOARD_HOST=0.0.0.0".into(),
            "-e".into(),
            "SPARK_MASTER_WEBUI_PORT=8082".into(),
            "-e".into(),
            "JUPYTER_DISABLE_TOKEN=1".into(),
            // Shrink the Spark master/worker JVM heaps (default 1g each) — the
            // biggest memory consumer on a node.
            "-e".into(),
            "SPARK_DAEMON_MEMORY=512m".into(),
        ]);
        // p2p mode: enable peer logic on the node. Nodes reach each other over
        // the host's WireGuard mesh, so no per-container tunnel (and no
        // NET_ADMIN/TUN device) is provisioned here any more.
        if p2p {
            args.extend([
                "-e".into(),
                "ZAKURO_P2P=true".into(),
                "-e".into(),
                format!("ZAKURO_PEER_KEY={}", peer_key()),
            ]);
        }
        args.push(image.clone());

        match docker(&args) {
            Ok(_) => println!("  • node {} up (dashboard :{})", name, dash),
            Err(e) => {
                eprintln!("  ✗ failed to start {}: {}", name, e);
                return;
            }
        }

        // Start the zakuro worker inside the node.
        let cmd = format!(
            "ZAKURO_WORKER_NAME={n} python -m zakuro.worker.server --host 0.0.0.0 --port {p} --worker-name {n} >/tmp/worker.log 2>&1",
            n = name,
            p = WORKER_PORT
        );
        let _ = d(&["exec", "-d", &name, "sh", "-c", &cmd]);
    }

    // Broker that discovers the nodes as a fleet.
    let _ = d(&["rm", "-f", BROKER]);
    let peers = (0..nodes)
        .map(|i| format!("{}:{}", node_name(i), WORKER_PORT))
        .collect::<Vec<_>>()
        .join(",");
    let mut bargs: Vec<String> = vec!["run".into(), "-d".into(), "--name".into(), BROKER.into()];
    if let Some(p) = &plat {
        bargs.push("--platform".into());
        bargs.push(p.clone());
    }
    bargs.extend([
        "--network".into(),
        NETWORK.into(),
        "--network-alias".into(),
        "broker".into(),
        "-p".into(),
        format!("{}:9000", BROKER_PORT),
        "-e".into(),
        format!("ZAKURO_PEERS={}", peers),
        "-e".into(),
        "ZAKURO_SCAN_INTERVAL=5".into(),
    ]);
    if p2p {
        bargs.extend([
            "-e".into(),
            "ZAKURO_P2P=true".into(),
            "-e".into(),
            format!("ZAKURO_PEER_KEY={}", peer_key()),
        ]);
    }
    bargs.extend([broker_image(), "broker".into()]);
    match docker(&bargs) {
        Ok(_) => println!("  • broker up (:{})", BROKER_PORT),
        Err(e) => {
            eprintln!("  ✗ failed to start broker: {}", e);
            return;
        }
    }

    // Wait for the broker to register the fleet.
    print!("  waiting for fleet");
    let _ = std::io::Write::flush(&mut std::io::stdout());
    let url = format!("http://localhost:{}/workers", BROKER_PORT);
    let mut found = 0;
    for _ in 0..30 {
        std::thread::sleep(Duration::from_secs(2));
        print!(".");
        let _ = std::io::Write::flush(&mut std::io::stdout());
        if let Ok(resp) = ureq::get(&url)
            .config()
            .timeout_global(Some(Duration::from_secs(3)))
            .build()
            .call()
        {
            if let Ok(body) = resp.into_body().read_to_string() {
                if let Ok(json) = serde_json::from_str::<serde_json::Value>(&body) {
                    found = json["total"].as_u64().unwrap_or(0) as usize;
                    if found >= nodes {
                        break;
                    }
                }
            }
        }
    }
    println!();
    println!(
        "  ✓ mesh up: {}/{} workers registered · dashboards :{}-:{} · broker :{}",
        found,
        nodes,
        DASH_BASE,
        DASH_BASE + (nodes as u16 - 1),
        BROKER_PORT
    );
    println!("  Run `zc workers` (or `/discovery` in the shell) to list the fleet.");
}

/// Tear down the mesh (nodes + broker). The network is left in place.
pub fn down() {
    println!("  zc mesh — tearing down");
    // Remove any zk<N> nodes plus the broker.
    let mut names: Vec<String> = (0..16).map(node_name).collect();
    names.push(BROKER.to_string());
    let mut removed = 0;
    for name in &names {
        if d(&["rm", "-f", name])
            .map(|s| !s.is_empty())
            .unwrap_or(false)
        {
            removed += 1;
        }
    }
    println!("  ✓ removed {} container(s)", removed);
}

/// Show the running mesh containers — name, short status, published host ports.
pub fn status() {
    // Pipe-delimited (no tabs) so it renders cleanly everywhere, then reformat.
    let out = d(&[
        "ps",
        "--filter",
        &format!("network={}", NETWORK),
        "--format",
        "{{.Names}}|{{.Status}}|{{.Ports}}",
    ]);
    let out = match out {
        Ok(o) => o,
        Err(e) => {
            eprintln!("  mesh status error: {}", e);
            return;
        }
    };
    if out.is_empty() {
        println!("  zc mesh — no nodes running (try `zc mesh up`)");
        return;
    }

    println!("  zc mesh — running containers:");
    println!("  {:<11} {:<10} PORTS", "NAME", "STATE");
    for line in out.lines() {
        let p: Vec<&str> = line.splitn(3, '|').collect();
        let name = p.first().copied().unwrap_or("-");
        // "Up 3 minutes (healthy)" -> "up" / "healthy".
        let raw_state = p.get(1).copied().unwrap_or("");
        let state = if raw_state.contains("healthy") {
            "healthy"
        } else if raw_state.starts_with("Up") {
            "up"
        } else {
            raw_state.split_whitespace().next().unwrap_or("-")
        };
        // Extract just the published host ports (e.g. "8090, 8265, …").
        let ports = p.get(2).copied().unwrap_or("");
        let mut host_ports: Vec<String> = ports
            .split(',')
            .filter_map(|seg| {
                let seg = seg.trim();
                seg.rsplit_once("->")
                    .and_then(|(left, _)| left.rsplit(':').next())
                    .map(|hp| hp.to_string())
            })
            .collect();
        host_ports.sort();
        host_ports.dedup();
        println!("  {:<11} {:<10} {}", name, state, host_ports.join(", "));
    }
}