datum_cli/

lib.rs

#![forbid(unsafe_code)]

use std::{
    collections::{HashMap, HashSet},
    fs,
    io::{self, Cursor, Write},
    net::SocketAddr,
    path::{Path, PathBuf},
    sync::Arc,
    time::Duration,
};

use clap::{Parser, Subcommand, error::ErrorKind};
use datum_agent::dcp::{
    ClientKind, DcpClient, DcpError, EventSubscription, Hello, MetricSubscription, ResponseStatus,
    proto::{
        ClusterJobList, ClusterJobNode, ClusterNodeError, ClusterNodeList, ClusterNodeStatus,
        Event, JobStatus as WireJobStatus, MetricSample, PlacementSpec, PlacementStrategy,
        StreamMetric,
    },
};
use datum_net::quic::{
    crypto::rustls::QuicClientConfig,
    quinn,
    rustls::{
        ClientConfig as RustlsClientConfig, RootCertStore,
        pki_types::{CertificateDer, PrivateKeyDer, PrivatePkcs8KeyDer},
    },
};
use serde::Serialize;

const DEFAULT_ADDR: &str = "127.0.0.1:9555";
const EXIT_OK: i32 = 0;
const EXIT_REMOTE: i32 = 1;
const EXIT_USAGE_OR_CONNECT: i32 = 2;
const EVENT_INITIAL_IDLE: Duration = Duration::from_millis(500);
const EVENT_IDLE: Duration = Duration::from_millis(250);
const METRIC_SNAPSHOT_WAIT: Duration = Duration::from_secs(2);
const DRAIN_POLL: Duration = Duration::from_millis(20);

#[derive(Debug, Parser)]
#[command(
    name = "datum",
    version,
    about = "Manage Datum jobs through the Datum Control Protocol"
)]
struct Cli {
    /// DCP address. Defaults to the loopback TCP development listener.
    #[arg(long, global = true, default_value = DEFAULT_ADDR)]
    addr: SocketAddr,

    /// CA certificate for QUIC+mTLS mode.
    #[arg(long, global = true, value_name = "PATH")]
    tls_ca: Option<PathBuf>,

    /// Client certificate for QUIC+mTLS mode.
    #[arg(long, global = true, value_name = "PATH")]
    tls_cert: Option<PathBuf>,

    /// Client private key for QUIC+mTLS mode.
    #[arg(long, global = true, value_name = "PATH")]
    tls_key: Option<PathBuf>,

    /// Emit machine-readable JSON instead of human tables.
    #[arg(long, global = true)]
    json: bool,

    #[command(subcommand)]
    command: Command,
}

#[derive(Debug, Subcommand)]
enum Command {
    /// List jobs.
    Ps {
        /// Aggregate jobs across cluster nodes.
        #[arg(long)]
        cluster: bool,
    },
    /// List cluster nodes.
    Nodes,
    /// Show one job.
    Status {
        #[arg(long)]
        cluster: bool,
        job: String,
    },
    /// Start a registered factory.
    Start {
        factory: String,
        #[arg(long)]
        name: Option<String>,
        #[arg(long = "param", value_name = "k=v")]
        params: Vec<String>,
    },
    /// Submit a job through cluster placement.
    Submit {
        #[arg(long)]
        cluster: bool,
        #[arg(long)]
        factory: String,
        #[arg(long)]
        name: Option<String>,
        #[arg(long = "param", value_name = "k=v")]
        params: Vec<String>,
        #[arg(long)]
        role: Option<String>,
        #[arg(long)]
        node: Option<String>,
    },
    /// Gracefully drain a job.
    Drain {
        #[arg(long)]
        cluster: bool,
        job: String,
    },
    /// Stop a job immediately.
    Stop {
        #[arg(long)]
        cluster: bool,
        job: String,
    },
    /// Restart a job.
    Restart { job: String },
    /// Tail lifecycle events.
    Events {
        #[arg(long)]
        follow: bool,
    },
    /// Show stream metrics for one job.
    Metrics {
        job: String,
        #[arg(long)]
        follow: bool,
    },
}

#[derive(Debug)]
enum CliError {
    Usage(String),
    Connect { addr: SocketAddr, message: String },
    Remote { message: String, fix: &'static str },
}

impl CliError {
    fn exit_code(&self) -> i32 {
        match self {
            Self::Remote { .. } => EXIT_REMOTE,
            Self::Usage(_) | Self::Connect { .. } => EXIT_USAGE_OR_CONNECT,
        }
    }

    fn print(&self) {
        match self {
            Self::Usage(message) => {
                eprintln!("{message}");
            }
            Self::Connect { addr, message } => {
                eprintln!("could not connect to datum-agent at {addr}: {message}");
                eprintln!(
                    "start datum-agent or pass --addr host:port; loopback TCP defaults to {DEFAULT_ADDR}"
                );
            }
            Self::Remote { message, fix } => {
                eprintln!("datum-agent rejected the request: {message}");
                eprintln!("{fix}");
            }
        }
    }
}

/// Run the CLI using process arguments and return the requested process exit code.
pub async fn main_entry() -> i32 {
    match Cli::try_parse() {
        Ok(cli) => match run(cli).await {
            Ok(()) => EXIT_OK,
            Err(error) => {
                error.print();
                error.exit_code()
            }
        },
        Err(error) => {
            let _ = error.print();
            clap_error_exit_code(&error)
        }
    }
}

/// clap prints `--help`/`--version`/missing-subcommand text as a "success" `Err`
/// (the message is already on stdout); only genuine usage errors should exit non-zero.
fn clap_error_exit_code(error: &clap::Error) -> i32 {
    match error.kind() {
        ErrorKind::DisplayHelp
        | ErrorKind::DisplayVersion
        | ErrorKind::DisplayHelpOnMissingArgumentOrSubcommand => EXIT_OK,
        _ => EXIT_USAGE_OR_CONNECT,
    }
}

async fn run(cli: Cli) -> Result<(), CliError> {
    validate_before_connect(&cli.command)?;
    let client = connect(&cli).await?;
    match &cli.command {
        Command::Ps { cluster } => {
            if *cluster {
                cluster_ps(&client, cli.addr, cli.json).await
            } else {
                ps(&client, cli.addr, cli.json).await
            }
        }
        Command::Nodes => nodes(&client, cli.addr, cli.json).await,
        Command::Status { cluster, job } => {
            status(&client, cli.addr, job, *cluster, cli.json).await
        }
        Command::Start {
            factory,
            name,
            params,
        } => {
            start(
                &client,
                cli.addr,
                factory,
                name.as_deref(),
                params,
                cli.json,
            )
            .await
        }
        Command::Submit {
            cluster,
            factory,
            name,
            params,
            role,
            node,
        } => {
            submit(
                &client,
                cli.addr,
                SubmitArgs {
                    cluster: *cluster,
                    factory,
                    name: name.as_deref(),
                    raw_params: params,
                    role: role.as_deref(),
                    node: node.as_deref(),
                },
                cli.json,
            )
            .await
        }
        Command::Drain { cluster, job } => drain(&client, cli.addr, job, *cluster, cli.json).await,
        Command::Stop { cluster, job } => stop(&client, cli.addr, job, *cluster, cli.json).await,
        Command::Restart { job } => restart(&client, cli.addr, job, cli.json).await,
        Command::Events { follow } => events(&client, cli.addr, *follow, cli.json).await,
        Command::Metrics { job, follow } => {
            metrics(&client, cli.addr, job, *follow, cli.json).await
        }
    }
}

fn validate_before_connect(command: &Command) -> Result<(), CliError> {
    match command {
        Command::Start { params, .. } => {
            parse_params(params)?;
        }
        Command::Submit {
            cluster,
            params,
            role,
            node,
            ..
        } => {
            if !cluster {
                return Err(CliError::Usage(
                    "`datum submit` currently requires --cluster; use `datum start` for local jobs."
                        .to_owned(),
                ));
            }
            if role.is_some() && node.is_some() {
                return Err(CliError::Usage(
                    "pass at most one of --role or --node for cluster placement.".to_owned(),
                ));
            }
            parse_params(params)?;
        }
        _ => {}
    }
    Ok(())
}

async fn connect(cli: &Cli) -> Result<DcpClient, CliError> {
    let hello = Hello::new(format!("datum-cli-{}", std::process::id()), ClientKind::Cli);
    match tls_paths(cli)? {
        Some(paths) => {
            let config = load_quic_client_config(paths.ca, paths.cert, paths.key)?;
            DcpClient::connect_quic(cli.addr, "localhost", config, hello)
                .await
                .map_err(|error| connect_error(cli.addr, error))
        }
        None => DcpClient::connect_tcp(cli.addr, hello)
            .await
            .map_err(|error| connect_error(cli.addr, error)),
    }
}

struct TlsPaths<'a> {
    ca: &'a Path,
    cert: &'a Path,
    key: &'a Path,
}

fn tls_paths(cli: &Cli) -> Result<Option<TlsPaths<'_>>, CliError> {
    match (&cli.tls_ca, &cli.tls_cert, &cli.tls_key) {
        (None, None, None) => Ok(None),
        (Some(ca), Some(cert), Some(key)) => Ok(Some(TlsPaths {
            ca: ca.as_path(),
            cert: cert.as_path(),
            key: key.as_path(),
        })),
        _ => Err(CliError::Usage(
            "QUIC+mTLS mode needs --tls-ca, --tls-cert, and --tls-key; pass all three or omit all three for loopback TCP.".to_owned(),
        )),
    }
}

fn load_quic_client_config(
    ca_path: &Path,
    cert_path: &Path,
    key_path: &Path,
) -> Result<quinn::ClientConfig, CliError> {
    let mut roots = RootCertStore::empty();
    for cert in read_certificates(ca_path)? {
        roots.add(cert).map_err(|error| {
            CliError::Usage(format!(
                "could not load --tls-ca {}: {error}; pass a PEM or DER CA certificate.",
                ca_path.display()
            ))
        })?;
    }
    let certs = read_certificates(cert_path)?;
    let key = read_private_key(key_path)?;
    let rustls = RustlsClientConfig::builder()
        .with_root_certificates(roots)
        .with_client_auth_cert(certs, key)
        .map_err(|error| {
            CliError::Usage(format!(
                "could not build QUIC client identity: {error}; check --tls-cert and --tls-key match."
            ))
        })?;
    let quic = QuicClientConfig::try_from(rustls).map_err(|error| {
        CliError::Usage(format!(
            "could not build QUIC client config: {error}; check the TLS files."
        ))
    })?;
    Ok(quinn::ClientConfig::new(Arc::new(quic)))
}

fn read_certificates(path: &Path) -> Result<Vec<CertificateDer<'static>>, CliError> {
    let bytes = fs::read(path).map_err(|error| {
        CliError::Usage(format!(
            "could not read {}: {error}; pass a readable PEM or DER certificate file.",
            path.display()
        ))
    })?;
    let mut cursor = Cursor::new(bytes.as_slice());
    let certs = rustls_pemfile::certs(&mut cursor)
        .collect::<Result<Vec<_>, _>>()
        .map_err(|error| {
            CliError::Usage(format!(
                "could not parse {}: {error}; pass a PEM or DER certificate file.",
                path.display()
            ))
        })?;
    if certs.is_empty() {
        Ok(vec![CertificateDer::from(bytes)])
    } else {
        Ok(certs)
    }
}

fn read_private_key(path: &Path) -> Result<PrivateKeyDer<'static>, CliError> {
    let bytes = fs::read(path).map_err(|error| {
        CliError::Usage(format!(
            "could not read {}: {error}; pass a readable PEM or DER PKCS#8 private key.",
            path.display()
        ))
    })?;
    let mut cursor = Cursor::new(bytes.as_slice());
    match rustls_pemfile::private_key(&mut cursor).map_err(|error| {
        CliError::Usage(format!(
            "could not parse {}: {error}; pass a PEM or DER PKCS#8 private key.",
            path.display()
        ))
    })? {
        Some(key) => Ok(key),
        None => Ok(PrivateKeyDer::Pkcs8(PrivatePkcs8KeyDer::from(bytes))),
    }
}

fn connect_error(addr: SocketAddr, error: DcpError) -> CliError {
    match error {
        DcpError::Response { status, message } => remote_error(status, message),
        other => CliError::Connect {
            addr,
            message: other.to_string(),
        },
    }
}

fn command_error(addr: SocketAddr, error: DcpError) -> CliError {
    match error {
        DcpError::Response { status, message } => remote_error(status, message),
        other => CliError::Connect {
            addr,
            message: other.to_string(),
        },
    }
}

fn remote_error(status: ResponseStatus, message: String) -> CliError {
    let fix = match status {
        ResponseStatus::BadRequest => "Fix the command arguments and try again.",
        ResponseStatus::Unauthorized => {
            "Check the DCP token or mTLS client certificate configured for datum-agent."
        }
        ResponseStatus::NotFound => "Check the job or factory name with `datum ps`.",
        ResponseStatus::Conflict => {
            "Pick a different job name, or stop the existing job before starting it again."
        }
        ResponseStatus::ProtocolMismatch => {
            "Upgrade the CLI and datum-agent to compatible versions."
        }
        ResponseStatus::DeadlineExceeded => {
            "Retry the command or check whether datum-agent is overloaded."
        }
        ResponseStatus::Ok | ResponseStatus::Failed => {
            "Check `datum status <job>` and the datum-agent logs, then retry."
        }
    };
    CliError::Remote { message, fix }
}

async fn ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
    let mut jobs = client
        .list_jobs()
        .await
        .map_err(|error| command_error(addr, error))?;
    jobs.sort_by(|left, right| left.name.cmp(&right.name));
    if json {
        print_json(&JobsOutput {
            jobs: jobs.iter().map(JobView::from).collect(),
        })?;
    } else if jobs.is_empty() {
        println!("no jobs");
    } else {
        print!("{}", render_jobs_table(&jobs));
    }
    Ok(())
}

async fn cluster_ps(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
    let mut list = client
        .list_cluster_jobs(1_000)
        .await
        .map_err(|error| command_error(addr, error))?;
    sort_cluster_jobs(&mut list);
    if json {
        print_json(&ClusterJobsOutput::from(&list))?;
    } else if list.nodes.iter().all(|node| node.jobs.is_empty()) && list.errors.is_empty() {
        println!("no cluster jobs");
    } else {
        print!("{}", render_cluster_jobs_table(&list));
        if !list.errors.is_empty() {
            print!("{}", render_cluster_errors_table(&list.errors));
        }
    }
    Ok(())
}

async fn nodes(client: &DcpClient, addr: SocketAddr, json: bool) -> Result<(), CliError> {
    let mut list = client
        .cluster_node_info(1_000)
        .await
        .map_err(|error| command_error(addr, error))?;
    list.nodes
        .sort_by(|left, right| left.node_id.cmp(&right.node_id));
    list.errors
        .sort_by(|left, right| left.node_id.cmp(&right.node_id));
    let coordinator_node_id = cluster_coordinator_node_id(client).await;
    if json {
        print_json(&ClusterNodesOutput::new(&list, coordinator_node_id.clone()))?;
    } else if list.nodes.is_empty() {
        println!("no cluster nodes");
    } else {
        print!(
            "{}",
            render_nodes_table(&list.nodes, coordinator_node_id.as_deref())
        );
        if !list.errors.is_empty() {
            print!("{}", render_cluster_errors_table(&list.errors));
        }
    }
    Ok(())
}

/// Best-effort placement-coordinator lookup: `ClusterNodeList` carries only membership,
/// so this reuses the `coordinator_node_id` already recorded on cluster jobs
/// (`WireJobStatus::coordinator_node_id`). Returns `None` when no cluster job has run yet.
async fn cluster_coordinator_node_id(client: &DcpClient) -> Option<String> {
    let jobs = client.list_cluster_jobs(1_000).await.ok()?;
    jobs.nodes
        .iter()
        .flat_map(|node| node.jobs.iter())
        .map(|job| job.coordinator_node_id.as_str())
        .find(|id| !id.is_empty())
        .map(str::to_owned)
}

async fn status(
    client: &DcpClient,
    addr: SocketAddr,
    job: &str,
    cluster: bool,
    json: bool,
) -> Result<(), CliError> {
    let status = if cluster {
        client.cluster_job_status(job, 1_000).await
    } else {
        client.job_status(job).await
    }
    .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&JobOutput {
            job: JobView::from(&status),
        })?;
    } else {
        print!("{}", render_status_table(&status));
    }
    Ok(())
}

struct SubmitArgs<'a> {
    cluster: bool,
    factory: &'a str,
    name: Option<&'a str>,
    raw_params: &'a [String],
    role: Option<&'a str>,
    node: Option<&'a str>,
}

async fn submit(
    client: &DcpClient,
    addr: SocketAddr,
    args: SubmitArgs<'_>,
    json: bool,
) -> Result<(), CliError> {
    if !args.cluster {
        return Err(CliError::Usage(
            "`datum submit` currently requires --cluster; use `datum start` for local jobs."
                .to_owned(),
        ));
    }
    let params = parse_params(args.raw_params)?;
    let instance_name = args.name.unwrap_or(args.factory);
    let placement = placement_spec(args.role, args.node);
    let status = client
        .submit_cluster_job(args.factory, instance_name, params, placement, 1_000)
        .await
        .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&ActionOutput {
            action: "submitted",
            message: format!("submitted cluster job '{}'", status.name),
            job: JobView::from(&status),
        })?;
    } else {
        println!(
            "submitted cluster job '{}' from factory '{}' to node '{}' — state {}, placement generation {}",
            status.name,
            args.factory,
            empty_dash(&status.placement_node_id),
            status.state,
            status.placement_generation
        );
    }
    Ok(())
}

fn placement_spec(role: Option<&str>, node: Option<&str>) -> PlacementSpec {
    PlacementSpec {
        role_constraint: role.unwrap_or_default().to_owned(),
        strategy: if node.is_some() {
            PlacementStrategy::Pinned as i32
        } else {
            PlacementStrategy::LeastJobs as i32
        },
        pinned_node_id: node.unwrap_or_default().to_owned(),
    }
}

async fn start(
    client: &DcpClient,
    addr: SocketAddr,
    factory: &str,
    name: Option<&str>,
    raw_params: &[String],
    json: bool,
) -> Result<(), CliError> {
    let params = parse_params(raw_params)?;
    let instance_name = name.unwrap_or(factory);
    let status = client
        .start_job(factory, instance_name, params)
        .await
        .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&ActionOutput {
            action: "started",
            message: format!("started '{}'", status.name),
            job: JobView::from(&status),
        })?;
    } else {
        println!(
            "started '{}' from factory '{}' — state {}, generation {}",
            status.name, factory, status.state, status.generation
        );
    }
    Ok(())
}

async fn drain(
    client: &DcpClient,
    addr: SocketAddr,
    job: &str,
    cluster: bool,
    json: bool,
) -> Result<(), CliError> {
    let initial = if cluster {
        client.drain_cluster_job(job, 1_000).await
    } else {
        client.drain_job(job).await
    }
    .map_err(|error| command_error(addr, error))?;
    let status = wait_for_drain(client, job, cluster, initial)
        .await
        .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&ActionOutput {
            action: "drained",
            message: format!("drained '{}'", status.name),
            job: JobView::from(&status),
        })?;
    } else if status.state == "Drained" {
        println!("drained '{}' — in-flight streams completed", status.name);
    } else {
        println!(
            "drain requested for '{}' — current state {}; run `datum status {}` to follow it",
            status.name, status.state, status.name
        );
    }
    Ok(())
}

async fn wait_for_drain(
    client: &DcpClient,
    job: &str,
    cluster: bool,
    mut status: WireJobStatus,
) -> Result<WireJobStatus, DcpError> {
    let wait = status
        .drain_remaining_ms
        .map(|ms| Duration::from_millis(ms.saturating_add(1_000)))
        .unwrap_or(Duration::from_secs(31))
        .min(Duration::from_secs(60));
    let started = tokio::time::Instant::now();
    while !is_drain_terminal(&status) && started.elapsed() < wait {
        tokio::time::sleep(DRAIN_POLL).await;
        status = if cluster {
            client.cluster_job_status(job, 1_000).await?
        } else {
            client.job_status(job).await?
        };
    }
    Ok(status)
}

fn is_drain_terminal(status: &WireJobStatus) -> bool {
    matches!(
        status.state.as_str(),
        "Completed" | "Drained" | "Stopped" | "Failed"
    )
}

async fn stop(
    client: &DcpClient,
    addr: SocketAddr,
    job: &str,
    cluster: bool,
    json: bool,
) -> Result<(), CliError> {
    let status = if cluster {
        client.stop_cluster_job(job, 1_000).await
    } else {
        client.stop_job(job).await
    }
    .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&ActionOutput {
            action: "stopped",
            message: format!("stopped '{}'", status.name),
            job: JobView::from(&status),
        })?;
    } else {
        println!("stopped '{}' — state {}", status.name, status.state);
    }
    Ok(())
}

async fn restart(
    client: &DcpClient,
    addr: SocketAddr,
    job: &str,
    json: bool,
) -> Result<(), CliError> {
    let status = client
        .restart_job(job)
        .await
        .map_err(|error| command_error(addr, error))?;
    if json {
        print_json(&ActionOutput {
            action: "restarted",
            message: format!("restarted '{}'", status.name),
            job: JobView::from(&status),
        })?;
    } else {
        println!(
            "restarted '{}' — state {}, generation {}",
            status.name, status.state, status.generation
        );
    }
    Ok(())
}

async fn events(
    client: &DcpClient,
    addr: SocketAddr,
    follow: bool,
    json: bool,
) -> Result<(), CliError> {
    let mut subscription = client
        .subscribe_events()
        .await
        .map_err(|error| command_error(addr, error))?;
    if follow {
        while let Some(event) = subscription.recv().await {
            print_event(&event, json)?;
        }
        if let Some(reason) = subscription.disconnect_reason() {
            return Err(CliError::Connect {
                addr,
                message: reason,
            });
        }
        return Ok(());
    }

    let events = collect_events(&mut subscription).await;
    if json {
        print_json(&EventsOutput {
            events: events.iter().map(EventView::from).collect(),
        })?;
    } else if events.is_empty() {
        println!("no events observed — run `datum events --follow` to keep waiting");
    } else {
        print!("{}", render_events_table(&events));
    }
    Ok(())
}

async fn collect_events(subscription: &mut EventSubscription) -> Vec<Event> {
    let mut events = Vec::new();
    loop {
        let timeout = if events.is_empty() {
            EVENT_INITIAL_IDLE
        } else {
            EVENT_IDLE
        };
        match tokio::time::timeout(timeout, subscription.recv()).await {
            Ok(Some(event)) => events.push(event),
            Ok(None) | Err(_) => break,
        }
    }
    events
}

fn print_event(event: &Event, json: bool) -> Result<(), CliError> {
    if json {
        print_json_line(&EventLine {
            event: EventView::from(event),
        })?;
    } else {
        println!(
            "{} {} '{}' generation {}",
            event.sequence, event.kind, event.name, event.generation
        );
        io::stdout().flush().map_err(io_usage_error)?;
    }
    Ok(())
}

async fn metrics(
    client: &DcpClient,
    addr: SocketAddr,
    job: &str,
    follow: bool,
    json: bool,
) -> Result<(), CliError> {
    let mut subscription = client
        .subscribe_metrics(0)
        .await
        .map_err(|error| command_error(addr, error))?;
    if follow {
        while let Some(sample) = subscription.recv().await {
            let filtered = filter_metric_sample(job, &sample);
            if filtered.streams.is_empty() {
                continue;
            }
            print_metric_sample_follow(job, &filtered, json)?;
        }
        if let Some(reason) = subscription.disconnect_reason() {
            return Err(CliError::Connect {
                addr,
                message: reason,
            });
        }
        return Ok(());
    }

    let sample = wait_for_metric_sample(job, &mut subscription).await;
    print_metric_sample(job, &sample, json)
}

async fn wait_for_metric_sample(
    job: &str,
    subscription: &mut MetricSubscription,
) -> MetricSampleView {
    let deadline = tokio::time::Instant::now() + METRIC_SNAPSHOT_WAIT;
    let mut latest = MetricSampleView {
        timestamp_ms: 0,
        streams: Vec::new(),
    };
    while tokio::time::Instant::now() < deadline {
        let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
        match tokio::time::timeout(
            remaining.min(Duration::from_millis(100)),
            subscription.recv(),
        )
        .await
        {
            Ok(Some(sample)) => {
                latest = filter_metric_sample(job, &sample);
                if !latest.streams.is_empty() {
                    break;
                }
            }
            Ok(None) => break,
            Err(_) => {}
        }
    }
    latest
}

fn print_metric_sample(job: &str, sample: &MetricSampleView, json: bool) -> Result<(), CliError> {
    if json {
        print_json(&MetricsOutput { job, sample })?;
    } else if sample.streams.is_empty() {
        println!("no metrics for '{job}' yet");
    } else {
        print!("{}", render_metrics_table(sample));
    }
    Ok(())
}

fn print_metric_sample_follow(
    job: &str,
    sample: &MetricSampleView,
    json: bool,
) -> Result<(), CliError> {
    if json {
        print_json_line(&MetricsOutput { job, sample })?;
    } else {
        print!("{}", render_metrics_table(sample));
        io::stdout().flush().map_err(io_usage_error)?;
    }
    Ok(())
}

fn filter_metric_sample(job: &str, sample: &MetricSample) -> MetricSampleView {
    let prefix = format!("{job}:");
    MetricSampleView {
        timestamp_ms: sample.timestamp_ms,
        streams: sample
            .streams
            .iter()
            .filter(|metric| metric.name == job || metric.name.starts_with(&prefix))
            .map(StreamMetricView::from)
            .collect(),
    }
}

fn sort_cluster_jobs(list: &mut ClusterJobList) {
    list.nodes
        .sort_by(|left, right| left.node_id.cmp(&right.node_id));
    for node in &mut list.nodes {
        node.jobs.sort_by(|left, right| left.name.cmp(&right.name));
    }
    list.errors
        .sort_by(|left, right| left.node_id.cmp(&right.node_id));
}

fn parse_params(raw_params: &[String]) -> Result<HashMap<String, String>, CliError> {
    let mut params = HashMap::new();
    let mut seen = HashSet::new();
    for raw in raw_params {
        let Some((key, value)) = raw.split_once('=') else {
            return Err(CliError::Usage(format!(
                "invalid --param '{raw}'; use --param key=value."
            )));
        };
        if key.trim().is_empty() {
            return Err(CliError::Usage(format!(
                "invalid --param '{raw}'; the key before '=' cannot be empty."
            )));
        }
        if !seen.insert(key.to_owned()) {
            return Err(CliError::Usage(format!(
                "duplicate --param key '{key}'; pass each key once."
            )));
        }
        params.insert(key.to_owned(), value.to_owned());
    }
    Ok(params)
}

#[derive(Serialize)]
struct JobsOutput {
    jobs: Vec<JobView>,
}

#[derive(Serialize)]
struct ClusterJobsOutput {
    partial: bool,
    nodes: Vec<ClusterJobNodeView>,
    errors: Vec<ClusterNodeErrorView>,
}

impl From<&ClusterJobList> for ClusterJobsOutput {
    fn from(list: &ClusterJobList) -> Self {
        Self {
            partial: list.partial,
            nodes: list.nodes.iter().map(ClusterJobNodeView::from).collect(),
            errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
        }
    }
}

#[derive(Serialize)]
struct ClusterNodesOutput {
    partial: bool,
    coordinator_node_id: Option<String>,
    nodes: Vec<ClusterNodeView>,
    errors: Vec<ClusterNodeErrorView>,
}

impl ClusterNodesOutput {
    fn new(list: &ClusterNodeList, coordinator_node_id: Option<String>) -> Self {
        Self {
            partial: list.partial,
            coordinator_node_id,
            nodes: list.nodes.iter().map(ClusterNodeView::from).collect(),
            errors: list.errors.iter().map(ClusterNodeErrorView::from).collect(),
        }
    }
}

#[derive(Serialize)]
struct JobOutput {
    job: JobView,
}

#[derive(Serialize)]
struct ActionOutput {
    action: &'static str,
    message: String,
    job: JobView,
}

#[derive(Serialize)]
struct EventsOutput {
    events: Vec<EventView>,
}

#[derive(Serialize)]
struct EventLine {
    event: EventView,
}

#[derive(Serialize)]
struct MetricsOutput<'a> {
    job: &'a str,
    sample: &'a MetricSampleView,
}

#[derive(Serialize)]
struct JobView {
    name: String,
    job_id: u64,
    state: String,
    desired_state: String,
    generation: u64,
    starts_total: u64,
    restarts_total: u64,
    last_start_at_ms: Option<u64>,
    last_exit_at_ms: Option<u64>,
    last_exit_reason: String,
    backoff_remaining_ms: Option<u64>,
    drain_remaining_ms: Option<u64>,
    drain_supported: bool,
    active_streams: Option<u64>,
    cluster_job: bool,
    factory_name: String,
    placement_strategy: String,
    role_constraint: String,
    coordinator_node_id: String,
    placement_node_id: String,
    placement_generation: u64,
    placement_history: Vec<PlacementHistoryView>,
}

impl From<&WireJobStatus> for JobView {
    fn from(status: &WireJobStatus) -> Self {
        Self {
            name: status.name.clone(),
            job_id: status.job_id,
            state: status.state.clone(),
            desired_state: status.desired_state.clone(),
            generation: status.generation,
            starts_total: status.starts_total,
            restarts_total: status.restarts_total,
            last_start_at_ms: status.last_start_at_ms,
            last_exit_at_ms: status.last_exit_at_ms,
            last_exit_reason: status.last_exit_reason.clone(),
            backoff_remaining_ms: status.backoff_remaining_ms,
            drain_remaining_ms: status.drain_remaining_ms,
            drain_supported: status.drain_supported,
            active_streams: status.active_streams,
            cluster_job: status.cluster_job,
            factory_name: status.factory_name.clone(),
            placement_strategy: placement_strategy_text(status),
            role_constraint: status
                .placement
                .as_ref()
                .map(|placement| placement.role_constraint.clone())
                .unwrap_or_default(),
            coordinator_node_id: status.coordinator_node_id.clone(),
            placement_node_id: status.placement_node_id.clone(),
            placement_generation: status.placement_generation,
            placement_history: status
                .placement_history
                .iter()
                .map(PlacementHistoryView::from)
                .collect(),
        }
    }
}

#[derive(Serialize)]
struct PlacementHistoryView {
    generation: u64,
    from_node_id: String,
    to_node_id: String,
    reason: String,
    timestamp_ms: u64,
}

impl From<&datum_agent::dcp::proto::ClusterPlacementHistory> for PlacementHistoryView {
    fn from(history: &datum_agent::dcp::proto::ClusterPlacementHistory) -> Self {
        Self {
            generation: history.generation,
            from_node_id: history.from_node_id.clone(),
            to_node_id: history.to_node_id.clone(),
            reason: history.reason.clone(),
            timestamp_ms: history.timestamp_ms,
        }
    }
}

fn placement_strategy_text(status: &WireJobStatus) -> String {
    let Some(placement) = status.placement.as_ref() else {
        return String::new();
    };
    match PlacementStrategy::try_from(placement.strategy).unwrap_or(PlacementStrategy::LeastJobs) {
        PlacementStrategy::LeastJobs => "LeastJobs".to_owned(),
        PlacementStrategy::Pinned => format!("Pinned({})", placement.pinned_node_id),
    }
}

#[derive(Serialize)]
struct ClusterJobNodeView {
    node_id: String,
    address: String,
    local: bool,
    jobs: Vec<JobView>,
}

impl From<&ClusterJobNode> for ClusterJobNodeView {
    fn from(node: &ClusterJobNode) -> Self {
        Self {
            node_id: node.node_id.clone(),
            address: node.address.clone(),
            local: node.local,
            jobs: node.jobs.iter().map(JobView::from).collect(),
        }
    }
}

#[derive(Serialize)]
struct ClusterNodeView {
    node_id: String,
    member_state: String,
    address: String,
    agent_addr: String,
    roles: Vec<String>,
    unreachable: bool,
    local: bool,
    session_state: String,
}

impl From<&ClusterNodeStatus> for ClusterNodeView {
    fn from(node: &ClusterNodeStatus) -> Self {
        Self {
            node_id: node.node_id.clone(),
            member_state: node.member_state.clone(),
            address: node.address.clone(),
            agent_addr: node.agent_addr.clone(),
            roles: node.roles.clone(),
            unreachable: node.unreachable,
            local: node.local,
            session_state: node.session_state.clone(),
        }
    }
}

#[derive(Serialize)]
struct ClusterNodeErrorView {
    node_id: String,
    message: String,
}

impl From<&ClusterNodeError> for ClusterNodeErrorView {
    fn from(error: &ClusterNodeError) -> Self {
        Self {
            node_id: error.node_id.clone(),
            message: error.message.clone(),
        }
    }
}

#[derive(Serialize)]
struct EventView {
    sequence: u64,
    timestamp_ms: u64,
    name: String,
    job_id: u64,
    generation: u64,
    kind: String,
    detail: String,
}

impl From<&Event> for EventView {
    fn from(event: &Event) -> Self {
        Self {
            sequence: event.sequence,
            timestamp_ms: event.timestamp_ms,
            name: event.name.clone(),
            job_id: event.job_id,
            generation: event.generation,
            kind: event.kind.clone(),
            detail: event.detail.clone(),
        }
    }
}

#[derive(Serialize)]
struct MetricSampleView {
    timestamp_ms: u64,
    streams: Vec<StreamMetricView>,
}

#[derive(Serialize)]
struct StreamMetricView {
    id: u64,
    name: String,
    elements_through: u64,
    restarts: u64,
    state: String,
    started_at_ms: u64,
    state_changed_at_ms: u64,
    finished_at_ms: Option<u64>,
    uptime_ms: u64,
}

impl From<&StreamMetric> for StreamMetricView {
    fn from(metric: &StreamMetric) -> Self {
        Self {
            id: metric.id,
            name: metric.name.clone(),
            elements_through: metric.elements_through,
            restarts: metric.restarts,
            state: metric.state.clone(),
            started_at_ms: metric.started_at_ms,
            state_changed_at_ms: metric.state_changed_at_ms,
            finished_at_ms: metric.finished_at_ms,
            uptime_ms: metric.uptime_ms,
        }
    }
}

fn print_json<T: Serialize>(value: &T) -> Result<(), CliError> {
    let json = serde_json::to_string_pretty(value).map_err(|error| {
        CliError::Usage(format!(
            "could not encode JSON output: {error}; retry without --json."
        ))
    })?;
    println!("{json}");
    Ok(())
}

fn print_json_line<T: Serialize>(value: &T) -> Result<(), CliError> {
    let json = serde_json::to_string(value).map_err(|error| {
        CliError::Usage(format!(
            "could not encode JSON output: {error}; retry without --json."
        ))
    })?;
    println!("{json}");
    io::stdout().flush().map_err(io_usage_error)?;
    Ok(())
}

fn io_usage_error(error: io::Error) -> CliError {
    CliError::Usage(format!(
        "could not write CLI output: {error}; check stdout/stderr."
    ))
}

fn render_jobs_table(jobs: &[WireJobStatus]) -> String {
    let rows = jobs
        .iter()
        .map(|job| {
            vec![
                job.name.clone(),
                job.job_id.to_string(),
                job.state.clone(),
                job.desired_state.clone(),
                job.generation.to_string(),
                format_number(job.starts_total),
                format_number(job.restarts_total),
                format_optional_number(job.active_streams),
            ]
        })
        .collect::<Vec<_>>();
    render_table(
        &[
            "JOB", "ID", "STATE", "DESIRED", "GEN", "STARTS", "RESTARTS", "ACTIVE",
        ],
        &rows,
    )
}

fn render_cluster_jobs_table(list: &ClusterJobList) -> String {
    let rows = list
        .nodes
        .iter()
        .flat_map(|node| {
            if node.jobs.is_empty() {
                vec![vec![
                    node.node_id.clone(),
                    node.address.clone(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                    "-".to_owned(),
                ]]
            } else {
                node.jobs
                    .iter()
                    .map(|job| {
                        vec![
                            node.node_id.clone(),
                            node.address.clone(),
                            job.name.clone(),
                            job.job_id.to_string(),
                            job.state.clone(),
                            job.desired_state.clone(),
                            job.generation.to_string(),
                            format_number(job.starts_total),
                            format_number(job.restarts_total),
                            cluster_marker(job),
                            empty_dash(&job.placement_node_id).to_owned(),
                            placement_generation_text(job),
                        ]
                    })
                    .collect::<Vec<_>>()
            }
        })
        .collect::<Vec<_>>();
    render_table(
        &[
            "NODE",
            "ADDRESS",
            "JOB",
            "ID",
            "STATE",
            "DESIRED",
            "GEN",
            "STARTS",
            "RESTARTS",
            "CLUSTER",
            "PLACED_ON",
            "PGEN",
        ],
        &rows,
    )
}

fn render_nodes_table(nodes: &[ClusterNodeStatus], coordinator_node_id: Option<&str>) -> String {
    let rows = nodes
        .iter()
        .map(|node| {
            vec![
                node.node_id.clone(),
                node.member_state.clone(),
                coordinator_marker(&node.node_id, coordinator_node_id).to_owned(),
                node.unreachable.to_string(),
                node.session_state.clone(),
                node.agent_addr.clone(),
                node.roles.join(","),
            ]
        })
        .collect::<Vec<_>>();
    render_table(
        &[
            "NODE",
            "MEMBER",
            "COORD",
            "UNREACHABLE",
            "SESSION",
            "AGENT_ADDR",
            "ROLES",
        ],
        &rows,
    )
}

fn coordinator_marker(node_id: &str, coordinator_node_id: Option<&str>) -> &'static str {
    if coordinator_node_id.is_some_and(|coordinator| coordinator == node_id) {
        "*"
    } else {
        "-"
    }
}

fn render_cluster_errors_table(errors: &[ClusterNodeError]) -> String {
    let rows = errors
        .iter()
        .map(|error| vec![error.node_id.clone(), error.message.clone()])
        .collect::<Vec<_>>();
    render_table(&["NODE", "ERROR"], &rows)
}

fn cluster_marker(job: &WireJobStatus) -> String {
    if job.cluster_job {
        "yes".to_owned()
    } else {
        "-".to_owned()
    }
}

fn placement_generation_text(job: &WireJobStatus) -> String {
    if job.cluster_job {
        job.placement_generation.to_string()
    } else {
        "-".to_owned()
    }
}

fn render_status_table(status: &WireJobStatus) -> String {
    let rows = vec![
        vec!["job".to_owned(), status.name.clone()],
        vec!["id".to_owned(), status.job_id.to_string()],
        vec!["state".to_owned(), status.state.clone()],
        vec!["desired".to_owned(), status.desired_state.clone()],
        vec!["generation".to_owned(), status.generation.to_string()],
        vec!["starts".to_owned(), format_number(status.starts_total)],
        vec!["restarts".to_owned(), format_number(status.restarts_total)],
        vec![
            "last_start_ms".to_owned(),
            format_optional_number(status.last_start_at_ms),
        ],
        vec![
            "last_exit_ms".to_owned(),
            format_optional_number(status.last_exit_at_ms),
        ],
        vec![
            "last_exit_reason".to_owned(),
            empty_dash(&status.last_exit_reason).to_owned(),
        ],
        vec![
            "backoff_remaining_ms".to_owned(),
            format_optional_number(status.backoff_remaining_ms),
        ],
        vec![
            "drain_remaining_ms".to_owned(),
            format_optional_number(status.drain_remaining_ms),
        ],
        vec![
            "drain_supported".to_owned(),
            status.drain_supported.to_string(),
        ],
        vec![
            "active_streams".to_owned(),
            format_optional_number(status.active_streams),
        ],
        vec!["cluster_job".to_owned(), status.cluster_job.to_string()],
        vec![
            "factory".to_owned(),
            empty_dash(&status.factory_name).to_owned(),
        ],
        vec![
            "placement_strategy".to_owned(),
            empty_dash(&placement_strategy_text(status)).to_owned(),
        ],
        vec![
            "role_constraint".to_owned(),
            status
                .placement
                .as_ref()
                .map(|placement| empty_dash(&placement.role_constraint).to_owned())
                .unwrap_or_else(|| "-".to_owned()),
        ],
        vec![
            "coordinator".to_owned(),
            empty_dash(&status.coordinator_node_id).to_owned(),
        ],
        vec![
            "placed_on".to_owned(),
            empty_dash(&status.placement_node_id).to_owned(),
        ],
        vec![
            "placement_generation".to_owned(),
            placement_generation_text(status),
        ],
    ];
    render_table(&["FIELD", "VALUE"], &rows)
}

fn render_events_table(events: &[Event]) -> String {
    let rows = events
        .iter()
        .map(|event| {
            vec![
                event.sequence.to_string(),
                event.timestamp_ms.to_string(),
                event.name.clone(),
                event.job_id.to_string(),
                event.generation.to_string(),
                event.kind.clone(),
                empty_dash(&event.detail).to_owned(),
            ]
        })
        .collect::<Vec<_>>();
    render_table(
        &["SEQ", "TIME_MS", "JOB", "ID", "GEN", "KIND", "DETAIL"],
        &rows,
    )
}

fn render_metrics_table(sample: &MetricSampleView) -> String {
    let rows = sample
        .streams
        .iter()
        .map(|metric| {
            vec![
                sample.timestamp_ms.to_string(),
                metric.name.clone(),
                metric.id.to_string(),
                metric.state.clone(),
                format_number(metric.elements_through),
                format_number(metric.restarts),
                metric.uptime_ms.to_string(),
            ]
        })
        .collect::<Vec<_>>();
    render_table(
        &[
            "TIME_MS",
            "STREAM",
            "ID",
            "STATE",
            "ELEMENTS",
            "RESTARTS",
            "UPTIME_MS",
        ],
        &rows,
    )
}

/// Render a simple left-aligned ASCII table with stable column widths.
#[must_use]
pub fn render_table(headers: &[&str], rows: &[Vec<String>]) -> String {
    let mut widths = headers
        .iter()
        .map(|header| header.len())
        .collect::<Vec<_>>();
    for row in rows {
        for (index, cell) in row.iter().enumerate().take(widths.len()) {
            widths[index] = widths[index].max(cell.len());
        }
    }

    let mut output = String::new();
    push_table_row(
        &mut output,
        &headers
            .iter()
            .map(|cell| (*cell).to_owned())
            .collect::<Vec<_>>(),
        &widths,
    );
    push_separator(&mut output, &widths);
    for row in rows {
        push_table_row(&mut output, row, &widths);
    }
    output
}

fn push_table_row(output: &mut String, row: &[String], widths: &[usize]) {
    for (index, width) in widths.iter().enumerate() {
        if index > 0 {
            output.push_str("  ");
        }
        let cell = row.get(index).map(String::as_str).unwrap_or("");
        output.push_str(cell);
        for _ in cell.len()..*width {
            output.push(' ');
        }
    }
    output.push('\n');
}

fn push_separator(output: &mut String, widths: &[usize]) {
    for (index, width) in widths.iter().enumerate() {
        if index > 0 {
            output.push_str("  ");
        }
        for _ in 0..*width {
            output.push('-');
        }
    }
    output.push('\n');
}

fn empty_dash(value: &str) -> &str {
    if value.is_empty() { "-" } else { value }
}

fn format_optional_number(value: Option<u64>) -> String {
    value.map(format_number).unwrap_or_else(|| "-".to_owned())
}

fn format_number(value: u64) -> String {
    let raw = value.to_string();
    let mut output = String::with_capacity(raw.len() + raw.len() / 3);
    for (index, ch) in raw.chars().rev().enumerate() {
        if index > 0 && index % 3 == 0 {
            output.push(',');
        }
        output.push(ch);
    }
    output.chars().rev().collect()
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn table_renderer_pads_columns_stably() {
        let output = render_table(
            &["JOB", "STATE", "STARTS"],
            &[
                vec!["ingest".to_owned(), "Running".to_owned(), "1".to_owned()],
                vec![
                    "daily-rollup".to_owned(),
                    "Drained".to_owned(),
                    "12,003".to_owned(),
                ],
            ],
        );

        assert_eq!(
            output,
            "JOB           STATE    STARTS\n\
             ------------  -------  ------\n\
             ingest        Running  1     \n\
             daily-rollup  Drained  12,003\n"
        );
    }

    #[test]
    fn params_require_key_value_pairs() {
        assert!(parse_params(&["threads=4".to_owned()]).is_ok());
        assert!(matches!(
            parse_params(&["threads".to_owned()]),
            Err(CliError::Usage(_))
        ));
        assert!(matches!(
            parse_params(&["threads=4".to_owned(), "threads=8".to_owned()]),
            Err(CliError::Usage(_))
        ));
    }

    #[test]
    fn numbers_use_operator_grouping() {
        assert_eq!(format_number(0), "0");
        assert_eq!(format_number(12_003), "12,003");
        assert_eq!(format_number(9_876_543_210), "9,876,543,210");
    }

    #[test]
    fn help_and_version_exit_zero() {
        let help_err = Cli::try_parse_from(["datum", "--help"]).unwrap_err();
        assert_eq!(clap_error_exit_code(&help_err), EXIT_OK);

        let subcommand_help_err = Cli::try_parse_from(["datum", "nodes", "--help"]).unwrap_err();
        assert_eq!(clap_error_exit_code(&subcommand_help_err), EXIT_OK);

        let version_err = Cli::try_parse_from(["datum", "--version"]).unwrap_err();
        assert_eq!(clap_error_exit_code(&version_err), EXIT_OK);
    }

    #[test]
    fn unknown_flag_exits_nonzero() {
        let error = Cli::try_parse_from(["datum", "--not-a-real-flag"]).unwrap_err();
        assert_ne!(clap_error_exit_code(&error), EXIT_OK);

        let missing_job = Cli::try_parse_from(["datum", "status"]).unwrap_err();
        assert_ne!(clap_error_exit_code(&missing_job), EXIT_OK);
    }

    fn cluster_node_status(node_id: &str) -> ClusterNodeStatus {
        ClusterNodeStatus {
            node_id: node_id.to_owned(),
            member_state: "Up".to_owned(),
            address: "127.0.0.1:9555".to_owned(),
            agent_addr: "127.0.0.1:9556".to_owned(),
            roles: vec!["worker".to_owned()],
            unreachable: false,
            local: false,
            session_state: "connected".to_owned(),
        }
    }

    #[test]
    fn coordinator_marker_flags_only_the_matching_node() {
        assert_eq!(coordinator_marker("node-a", Some("node-a")), "*");
        assert_eq!(coordinator_marker("node-a", Some("node-b")), "-");
        assert_eq!(coordinator_marker("node-a", None), "-");
    }

    #[test]
    fn nodes_table_marks_the_coordinator_row() {
        let nodes = vec![cluster_node_status("node-a"), cluster_node_status("node-b")];
        let output = render_nodes_table(&nodes, Some("node-b"));

        assert!(output.contains("COORD"));
        let lines: Vec<&str> = output.lines().collect();
        let node_a_row = lines[2];
        let node_b_row = lines[3];
        assert!(!node_a_row.split_whitespace().any(|token| token == "*"));
        assert!(node_b_row.split_whitespace().any(|token| token == "*"));
    }

    #[test]
    fn cluster_nodes_output_json_carries_coordinator_node_id() {
        let list = ClusterNodeList {
            nodes: vec![cluster_node_status("node-a")],
            errors: Vec::new(),
            partial: false,
        };
        let output = ClusterNodesOutput::new(&list, Some("node-a".to_owned()));

        let value = serde_json::to_value(&output).expect("serializes");
        assert_eq!(value["coordinator_node_id"], "node-a");
        assert_eq!(value["nodes"][0]["node_id"], "node-a");
    }
}