irtt-rs/crates/irtt-app/src/cmd/client/args.rs

use std::time::Duration;

use crate::shared::client::{
    parse_target, parse_test_duration, prepare_managed_run, CommonClientArgs, GroupPacingArg,
    ManagedRunSetup, TargetArg, TargetSelection, TimestampArg,
};
use clap::{Parser, ValueEnum};

pub const DEFAULT_CLIENT_DURATION: Duration = Duration::from_secs(10);

#[derive(Debug, Clone, Parser)]
#[command(name = "irtt-client", about = "Minimal IRTT-compatible stream client")]
pub struct ClientArgs {
    /// Server address/host, optionally prefixed with LABEL= and suffixed with @hmac=KEY. Repeat for multi-target mode.
    #[arg(
        value_name = "TARGET",
        num_args = 0..,
        value_parser = parse_target,
        long_help = "Server address/host, optionally prefixed with LABEL= and suffixed with @hmac=KEY. Repeat for multi-target mode.\n\nExamples:\n  irtt-client host-a:2112 host-b:2112\n  irtt-client eu=host-a:2112 us=host-b:2112\n  irtt-client ams=host-a:2112@hmac=secret public=host-b:2112@hmac=\n\nA target without @hmac= inherits --hmac. @hmac= explicitly disables it for that target."
    )]
    pub targets: Vec<TargetArg>,

    /// Read complete desired target sets from standard input in continuous mode.
    #[arg(
        long,
        long_help = "Read one complete desired target set per stdin line in continuous mode. Commas separate targets and backslash escapes commas within a target. [] selects an empty desired set. Records are limited to 64 KiB. EOF stops the client gracefully and may discard a revision that has not yet been acknowledged as applied."
    )]
    pub targets_stdin: bool,

    /// Managed group pacing for multi-target mode.
    #[arg(
        long,
        value_enum,
        default_value_t = GroupPacingArg::Staggered,
        long_help = "Managed group pacing for multi-target mode.\n\nstaggered spaces active targets across the probe interval. burst sends one probe to every active target back-to-back once per interval."
    )]
    pub pacing: GroupPacingArg,

    #[arg(
        long,
        default_value = "10s",
        value_parser = parse_test_duration,
        help = "Test duration; use 0 for continuous mode",
        long_help = "Test duration; use 0 for continuous mode.\n\nFinite runs retain exact statistics for final summaries. Continuous mode uses bounded-memory running statistics and prints a final summary when interrupted; --targets-stdin also prints the final retained summary on EOF."
    )]
    pub duration: Duration,

    #[command(flatten)]
    pub common: CommonClientArgs,

    /// Event row output format.
    #[arg(
        long,
        value_enum,
        default_value_t = OutputFormat::Table,
        long_help = "Event row output format.\n\nTable is the default interactive format. CSV, TSV, and JSON Lines default to all columns for structured export."
    )]
    pub format: OutputFormat,

    /// Comma-separated event row columns, or default/all.
    #[arg(
        short = 'c',
        long,
        value_name = "COLUMNS",
        long_help = "Comma-separated event row columns, or default/all.\n\nThe default table columns are compact and hide echo_sent rows. Custom table columns include all event rows. Run --list-columns to see valid names."
    )]
    pub columns: Option<String>,

    /// List available event row columns and aliases, then exit.
    #[arg(
        long,
        long_help = "List available event row columns and aliases, then exit.\n\nA server argument is not required when listing columns."
    )]
    pub list_columns: bool,

    /// Header policy for table, CSV, and TSV output.
    #[arg(
        long,
        value_enum,
        default_value_t = HeaderMode::Auto,
        long_help = "Header policy for table, CSV, and TSV output.\n\nJSON Lines never prints a header."
    )]
    pub header: HeaderMode,

    /// Include extra fields in table output and final summaries.
    #[arg(long)]
    pub verbose: bool,
}

impl ClientArgs {
    /// Validate the selected targets and prepare the managed run.
    ///
    /// Targets are only required here. `--list-columns` returns before this is
    /// called, so listing columns still needs no target.
    pub fn prepare(&self) -> Result<ManagedRunSetup, String> {
        if self.targets_stdin && !self.is_continuous() {
            return Err("--targets-stdin requires --duration 0".to_owned());
        }
        prepare_managed_run(
            &self.common,
            self.duration,
            TargetSelection {
                targets: &self.targets,
                pacing: self.pacing,
                stdin_controlled: self.targets_stdin,
            },
        )
    }

    pub fn is_continuous(&self) -> bool {
        self.duration == Duration::ZERO
    }

    pub fn timestamp_mode(&self) -> TimestampArg {
        self.common.tstamp
    }
}

impl std::ops::Deref for ClientArgs {
    type Target = CommonClientArgs;

    fn deref(&self) -> &Self::Target {
        &self.common
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum OutputFormat {
    /// Readable terminal table output.
    Table,
    /// Comma-separated event rows.
    Csv,
    /// Tab-separated event rows.
    Tsv,
    /// One JSON object per event row.
    Jsonl,
}

impl OutputFormat {
    pub fn prints_summary(self) -> bool {
        matches!(self, Self::Table)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum HeaderMode {
    /// Print headers for table, CSV, and TSV output.
    Auto,
    /// Always print headers where the format supports them.
    Always,
    /// Never print headers.
    Never,
}

#[cfg(test)]
mod tests {
    use super::*;
    use clap::{CommandFactory, Parser};
    use irtt_client::NegotiationPolicy;
    use irtt_proto::{Clock, ReceivedStats, StampAt};

    fn parse(args: &[&str]) -> Result<ClientArgs, clap::Error> {
        let mut argv = vec!["irtt-client"];
        argv.extend_from_slice(args);
        ClientArgs::try_parse_from(argv)
    }

    #[test]
    fn client_parser_keeps_finite_default_duration() {
        let args = parse(&["127.0.0.1:2112"]).unwrap();

        assert_eq!(args.targets.len(), 1);
        assert_eq!(args.prepare().unwrap().targets[0].label, "127.0.0.1:2112");
        assert_eq!(args.pacing, GroupPacingArg::Staggered);
        assert_eq!(args.duration, DEFAULT_CLIENT_DURATION);
        assert_eq!(args.interval, Duration::from_secs(1));
        assert_eq!(args.format, OutputFormat::Table);
        assert_eq!(args.columns, None);
        assert_eq!(args.header, HeaderMode::Auto);
        assert!(!args.targets_stdin);
        assert!(!args.is_continuous());
        assert_eq!(args.timestamp_mode(), TimestampArg::Both);

        let config = args.prepare().unwrap().client;
        assert_eq!(config.duration, Some(DEFAULT_CLIENT_DURATION));
        assert_eq!(config.received_stats, ReceivedStats::Both);
        assert_eq!(config.stamp_at, StampAt::Both);
        assert_eq!(config.clock, Clock::Both);
        assert_eq!(config.negotiation_policy, NegotiationPolicy::Strict);
    }

    #[test]
    fn client_parser_accepts_stream_formats_and_rejects_tui_output() {
        for format in ["table", "csv", "tsv", "jsonl"] {
            assert!(parse(&["--format", format, "127.0.0.1:2112"]).is_ok());
        }
        assert!(parse(&["--format", "tui", "127.0.0.1:2112"]).is_err());
    }

    #[test]
    fn list_columns_does_not_require_server() {
        let args = parse(&["--list-columns"]).unwrap();
        assert!(args.targets.is_empty());
        assert!(args.list_columns);
        // The parser accepts no target; preparation is where a target becomes
        // mandatory, and the stream command returns before preparing when
        // listing columns.
        let err = args.prepare().unwrap_err();
        assert!(err.contains("--list-columns"));
    }

    #[test]
    fn multiple_positional_targets_parse_as_targets() {
        let args = parse(&["host-a:2112", "host-b:2112"]).unwrap();
        let specs = args.prepare().unwrap().targets;

        assert_eq!(specs[0].label, "host-a:2112");
        assert_eq!(specs[0].managed.server_addr, "host-a:2112");
        assert_eq!(specs[1].label, "host-b:2112");
        assert_eq!(specs[1].managed.server_addr, "host-b:2112");
    }

    #[test]
    fn repeated_labeled_targets_parse() {
        let args = parse(&["ams=ams.example.com:2112", "sg=sg.example.com:2112"]).unwrap();
        let specs = args.prepare().unwrap().targets;

        assert_eq!(specs[0].label, "ams");
        assert_eq!(specs[0].managed.server_addr, "ams.example.com:2112");
        assert_eq!(specs[1].label, "sg");
        assert_eq!(specs[1].managed.server_addr, "sg.example.com:2112");
    }

    #[test]
    fn mixed_labeled_and_unlabeled_targets_preserve_argument_order() {
        let args = parse(&["local.example", "eu=a.example", "backup.example"]).unwrap();
        let specs = args.prepare().unwrap().targets;

        let labels: Vec<_> = specs.iter().map(|target| target.label.as_str()).collect();
        assert_eq!(labels, ["local.example", "eu", "backup.example"]);
    }

    #[test]
    fn duplicate_labels_are_rejected() {
        let args = parse(&["host-a:2112", "host-a:2112=host-b:2112"]).unwrap();
        let err = args.prepare().unwrap_err();

        assert!(err.contains("duplicate target label"));
    }

    #[test]
    fn duplicate_positional_target_strings_get_stable_suffixes() {
        let args = parse(&["host-a:2112", "host-a:2112"]).unwrap();
        let specs = args.prepare().unwrap().targets;

        assert_eq!(specs[0].label, "host-a:2112");
        assert_eq!(specs[1].label, "host-a:2112#2");
    }

    #[test]
    fn duplicate_target_endpoints_are_allowed() {
        let args = parse(&["127.0.0.1:2112", "127.0.0.1"]).unwrap();
        assert_eq!(args.prepare().unwrap().target_count(), 2);
    }

    #[test]
    fn invalid_labelled_target_syntax_is_rejected() {
        assert!(parse(&["=127.0.0.1:2112"]).unwrap().prepare().is_err());
        assert!(parse(&["label="]).unwrap().prepare().is_err());
    }

    #[test]
    fn old_target_option_is_rejected() {
        assert!(parse(&["--target", "eu=host.example"]).is_err());
    }

    #[test]
    fn pacing_option_accepts_supported_values() {
        assert_eq!(
            parse(&["--pacing", "staggered", "127.0.0.1:2112"])
                .unwrap()
                .pacing,
            GroupPacingArg::Staggered
        );
        assert_eq!(
            parse(&["--pacing", "burst", "127.0.0.1:2112"])
                .unwrap()
                .pacing,
            GroupPacingArg::Burst
        );
    }

    #[test]
    fn client_help_lists_shared_protocol_options() {
        let help = ClientArgs::command().render_help().to_string();
        assert!(help.contains("--format <FORMAT>"));
        assert!(!help.contains("--target "));
        assert!(help.contains("--pacing <PACING>"));
        assert!(help.contains("--columns <COLUMNS>"));
        assert!(help.contains("--list-columns"));
        assert!(help.contains("--header <HEADER>"));
        assert!(help.contains("--tstamp <MODE>"));
        assert!(help.contains("--stats <STATS>"));
        assert!(help.contains("--sfill <STRING>"));
        assert!(help.contains("--dscp <0..=63>"));
        assert!(help.contains("--ttl <1..=255>"));
        assert!(help.contains("--loose"));
    }

    #[test]
    fn an_empty_hmac_key_is_rejected() {
        let err = parse(&["--hmac", "", "127.0.0.1:2112"]).unwrap_err();
        assert!(err.to_string().contains("HMAC key must not be empty"));
    }

    #[test]
    fn a_non_empty_hmac_key_is_accepted() {
        let args = parse(&["--hmac", "secret", "127.0.0.1:2112"]).unwrap();
        let config = args.prepare().unwrap().client;
        assert_eq!(config.hmac_key, Some(b"secret".to_vec()));
    }

    #[test]
    fn targets_stdin_requires_continuous_mode_and_allows_empty_initial_set() {
        assert!(parse(&["--targets-stdin"]).unwrap().prepare().is_err());
        let args = parse(&["--duration", "0", "--targets-stdin"]).unwrap();
        let setup = args.prepare().unwrap();
        assert!(setup.targets.is_empty());
        assert!(setup.stdin_controlled);
        let config = setup.managed_config();
        assert_eq!(
            config.completion,
            irtt_client::managed::ManagedCompletionPolicy::ExplicitStop
        );
        assert_eq!(config.max_live_target_generations, 256);
    }

    #[test]
    fn output_format_summary_policy_is_table_only() {
        assert!(OutputFormat::Table.prints_summary());
        assert!(!OutputFormat::Csv.prints_summary());
        assert!(!OutputFormat::Tsv.prints_summary());
        assert!(!OutputFormat::Jsonl.prints_summary());
    }
}