#compdef iperf_orchestrator.sh
# zsh completion for iperf_orchestrator.sh
#
# Install: copy to a directory on $fpath (e.g. /usr/local/share/zsh/site-functions/_iperf_orchestrator)
# then `compinit` (or restart your shell).

_iperf_orchestrator() {
    local -a subcommands run_modes global_flags
    subcommands=(
        'gen:Write the plan file (hosts + settings in one artifact)'
        'start:Start iperf2 daemons everywhere and run the tests'
        'summarize:Collect results, render CSV/pivot/heatmap, print hints'
        'stop:Stop the iperf2 daemons (logs stay on the hosts)'
        'clean:Stop + delete the remote work dir everywhere, verified'
        'run:start + summarize + stop + clean in one shot'
        'hints:What you want to know -> the command for it'
        'status:Probe hosts, list runs, show live progress'
        'check-iperf:Probe iperf2 + mpstat on every host'
        'check-servers:Probe for running iperf daemons'
        'start-servers:Start iperf -s on every host'
        'create-scripts:Generate per-host run scripts locally'
        'distribute-scripts:Copy run scripts to each host'
        'run-tests:Run the iperf2 mesh tests'
        'collect-results:Pull all per-host logs back'
        'stop-servers:Stop iperf -s daemons'
        'cleanup:Remove $REMOTE_DIR on every host (requires --yes)'
        'parse-csv:Parse iperf2 logs into CSV'
        'parse-cpu:Parse mpstat / fallback CPU samples'
        'make-pivot:Render text pivot table'
        'make-heatmap:Render heatmap PNG'
        'doctor:Probe local prerequisites'
        'results-summary:Print P50/P95/slowest pairs from CSV'
        'export-overlay:Write the run as datacenter layout viewer overlay samples'
        'all:Run the full pipeline end-to-end'
        'help:Show the common commands and flags'
        'help-advanced:Every command, every flag, every env var'
        'version:Print the version and exit'
    )

    run_modes=(parallel sequential-host sequential-pair rolling)

    global_flags=(
        '--plan[plan file written by gen]:file:_files'
        '--servers[server list file]:file:_files'
        '-s[server list file]:file:_files'
        '--output[results base directory]:dir:_files -/'
        '-o[results base directory]:dir:_files -/'
        '--run-id[address an existing run]:run-id:'
        '--port[iperf2 listening port]:port:'
        '--duration[seconds per pair]:seconds:'
        '-d[seconds per pair]:seconds:'
        '--test-timeout[hard cap per iperf -c (default duration+30; 0 = off)]:seconds:'
        '--single-server[sequential-host: one all->one round against HOST]:host:'
        '--streams[parallel TCP streams per test]:N:'
        '-P[parallel TCP streams per test]:N:'
        '--ssh-jobs[max concurrent SSH/SCP fan-out]:N:'
        '-j[max concurrent SSH/SCP fan-out]:N:'
        '--start-delay[synchronized-start lead time]:seconds:'
        '--total-time[rolling mode wall-time]:seconds:'
        '--host-flows[concurrent iperf procs per directed edge]:N:'
        '--bandwidth[cap per-flow rate (iperf2 -b)]:rate:'
        '-b[cap per-flow rate (iperf2 -b)]:rate:'
        '--length[TCP read/write buffer size (iperf2 -l)]:size:'
        '-l[TCP read/write buffer size (iperf2 -l)]:size:'
        '--window[TCP window / socket buffer (iperf2 -w)]:size:'
        '-w[TCP window / socket buffer (iperf2 -w)]:size:'
        '--mss[TCP maximum segment size (iperf2 -M)]:size:'
        '-M[TCP maximum segment size (iperf2 -M)]:size:'
        '--no-nagle[disable Nagle'"'"'s algorithm (iperf2 -N)]'
        '-N[disable Nagle'"'"'s algorithm (iperf2 -N)]'
        '--bind[bind clients to matching NIC (substring pattern)]:pattern:'
        '-B[bind clients to matching NIC (substring pattern)]:pattern:'
        '--server-bind[bind the daemon side too]:pattern:'
        '--ssh-user[SSH login user]:user:'
        '-u[SSH login user]:user:'
        '--remote-dir[remote working dir]:dir:'
        '--python[Python interpreter]:path:_files'
        '--overlay[also export the overlay during process]'
        '--overlay-out[overlay destination (- for stdout)]:file:_files'
        '--overlay-format[overlay format]:format:(tsv ndjson)'
        '--overlay-map[<host> <element> mapping file]:file:_files'
        '--overlay-prefix[string prepended to every target]:prefix:'
        '--overlay-append[append instead of replacing the destination]'
        '--overlay-reduce[one sample per host per test (median)]'
        '--overlay-no-meta[omit the !test metadata lines]'
        '--dry-run[print SSH/SCP commands without executing]'
        '-n[print SSH/SCP commands without executing]'
        '--verbose[also print every ssh/scp invocation]'
        '-v[also print every ssh/scp invocation]'
        '--quiet[suppress non-WARN/ERROR log lines]'
        '-q[suppress non-WARN/ERROR log lines]'
        '-h[show help]'
        '--help[show help]'
        '--help-advanced[show every command, flag, and env var]'
        '--version[print the version and exit]'
    )

    _arguments -C \
        $global_flags \
        '1: :->subcmd' \
        '*::: :->subargs'

    case $state in
        subcmd)
            _describe -t subcommands 'subcommand' subcommands
            ;;
        subargs)
            case $line[1] in
                run-tests|all)
                    _arguments \
                        "1:mode:(${run_modes[*]})" \
                        '--keep-going[continue past per-host failures]' \
                        '--help[show help]'
                    ;;
                gen)
                    _arguments \
                        "1:mode:(${run_modes[*]})" \
                        '--grid[write the hosts as a src\\dst pair grid]' \
                        '--help[show help]'
                    ;;
                start)
                    _arguments \
                        "1:mode:(${run_modes[*]})" \
                        '--keep-going[continue past per-host failures]' \
                        '--help[show help]'
                    ;;
                run)
                    _arguments \
                        "1:mode:(${run_modes[*]})" \
                        '--for[total-time (rolling) or per-test duration]:seconds:' \
                        '--keep-going[continue past per-host failures]' \
                        '--help[show help]'
                    ;;
                cleanup)
                    _arguments '--yes[confirm rm -rf on every host]' '--help[show help]'
                    ;;
                status)
                    _arguments '--watch[redraw every N seconds]:seconds:' '--help[show help]'
                    ;;
                *)
                    _arguments '--help[show help]'
                    ;;
            esac
            ;;
    esac
}

_iperf_orchestrator "$@"
