# Agent Workbench CLI - Zsh Completion
#
# Install: place in a directory in your $fpath
#   mv _awb /usr/local/share/zsh/site-functions/
#
# Or source from ~/.zshrc:
#   fpath=(/path/to/completions $fpath)
#   autoload -Uz compinit && compinit

#compdef awb

_awb() {
    local -a commands
    commands=(
        'init:Scaffold workbench files in the current repository'
        'learn:Inspect the current repository and generate discovery artifacts'
        'brief:Create, finalize, or inspect behavior briefs'
        'evals:Generate, review, and approve eval suites'
        'run:Execute an agent against an eval subset'
        'run-chain:Compatibility alias for awb run chain'
        'run-workflow:Compatibility alias for awb run workflow'
        'compare:Compare an experiment against a baseline or the best comparable run'
        'check:Run loop-readiness checks for Codex optimization loops'
        'status:Show the current optimization status for one or more streams'
        'next:Collapse status, compare, diagnose, and proposal guidance into one report'
        'diagnose:Analyze the current run and suggest the next mutation classes'
        'propose:Generate bounded mutation proposals from a diagnosis'
        'iterate:Turn diagnosis and proposals into measured optimization lanes'
        'optimize:Run a bounded optimize loop with search and gates'
        'ablate:Run one-surface ablations against an agent or experiment'
        'matrix:Run a cross-agent or cross-provider matrix on the same eval subset'
        'factorial:Run small factorial experiments across prompts, tools, or MCPs'
        'path-compare:Summarize dialogue path behavior from a completed experiment'
        'patch-template:Turn a proposal into an editable implementation template'
        'promote:Mark a candidate as the promoted baseline for its stream'
        'release:Summarize release readiness across benchmark streams'
        'adapters:List available builtin and registered adapters'
        'benchmarks:Register, list, and install portable benchmark bundles'
        'redteam:Generate adversarial suites and summarize red-team runs'
        'human-review:Manage human review queues and answers'
        'artifacts:Inspect and export experiment artifacts'
        'scorer:Author and lint custom scorer files'
        'scorer-audit:Audit scorer judgment metadata and calibration'
        'route-policy:Author, lint, and review dialogue route policies'
        'migrate:Rewrite stored artifacts to the current schema version'
        'doctor:Check workbench installation health and configuration'
        'traces:Validate and explain trace telemetry from experiments'
    )

    local -a format_args
    format_args=(
        '--format[Output format]:format:(text json)'
    )

    _arguments -C \
        '--version[Show version]' \
        '--help[Show help]' \
        '1: :->command' \
        '*:: :->args'

    case $state in
        command)
            _describe 'command' commands
            ;;
        args)
            case $words[2] in
                brief)
                    _arguments \
                        '1: :->brief_command'
                    case $words[3] in
                        start)
                            _arguments \
                                '--brief-id[Stable identifier for the behavior brief]:brief_id' \
                                '--goal[High-level goal to seed the draft brief]:goal' \
                                '--agent-id[Agent spec to use]:agent_id'
                            ;;
                        draft)
                            _arguments \
                                '--brief-id[Behavior brief identifier]:brief_id' \
                                '--agent-id[Agent spec to use]:agent_id' \
                                '--goal[Optional goal override]:goal'
                            ;;
                        finalize)
                            _arguments \
                                '--brief-id[Behavior brief identifier]:brief_id' \
                                '--agent-id[Agent spec to use]:agent_id'
                            ;;
                        packet)
                            _arguments \
                                '--brief-id[Behavior brief identifier]:brief_id' \
                                '--agent-id[Agent spec to use]:agent_id'
                            ;;
                    esac
                    ;;
                evals)
                    _arguments \
                        '1: :->evals_command'
                    case $words[3] in
                        generate)
                            _arguments \
                                '--brief-id[Approved behavior brief identifier]:brief_id' \
                                '--eval-id[Stable identifier for the generated eval suite]:eval_id'
                            ;;
                        expand)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--brief-id[Optional approved brief override]:brief_id' \
                                '--source[Expansion source]:source' \
                                '--brief-limit[Maximum brief-derived cases]:limit' \
                                '--experiment[Experiment id]:experiment' \
                                '--agent-id[Agent id]:agent_id' \
                                '--subset[Subset]:subset' \
                                '--benchmark-signature[Benchmark signature]:benchmark_signature' \
                                '--case-id[Case id]:case_id' \
                                '--tag[Tag]:tag' \
                                '--match[Match text]:text' \
                                '--cluster[Failure cluster id]:cluster' \
                                '--focus-limit[Maximum failure-driven cases]:limit'
                            ;;
                        review|approve)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id'
                            ;;
                        export)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--output[Optional JSON output path]:output' \
                                '--benchmark-id[Optional benchmark identifier]:benchmark_id'
                            ;;
                        import)
                            _arguments \
                                '--input[Path to a benchmark bundle JSON file]:input:_files' \
                                '--eval-id[Optional eval id override]:eval_id' \
                                '--force[Overwrite existing eval suite]'
                            ;;
                    esac
                    ;;
                run)
                    _arguments \
                        '1: :->run_command'
                    case $words[3] in
                        chain)
                            _arguments \
                                '--agent-id[Chain step agent id, repeat per step]:agent_id' \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--subset[Named eval split to run]:subset' \
                                '--chain-handoff-template[Input template for steps after the first]:template' \
                                '--case-id[Restrict to specific case id]:case_id' \
                                '--tag[Restrict to cases with given tag]:tag' \
                                '--match[Restrict to cases matching text]:text' \
                                '--case-limit[Maximum number of matching cases]:limit'
                            ;;
                        workflow)
                            _arguments \
                                '--workflow-file[Workflow YAML/JSON path]:workflow_file' \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--subset[Named eval split to run]:subset' \
                                '--workflow-handoff-template[Input template for handoffs]:template' \
                                '--case-id[Restrict to specific case id]:case_id' \
                                '--tag[Restrict to cases with given tag]:tag' \
                                '--match[Restrict to cases matching text]:text' \
                                '--case-limit[Maximum number of matching cases]:limit'
                            ;;
                        replay)
                            _arguments \
                                '--experiment[Source experiment id]:experiment' \
                                '--case-id[Case id to replay]:case_id' \
                                '--agent-id[Override agent id]:agent_id' \
                                '--eval-id[Override eval id]:eval_id' \
                                '--subset[Override subset]:subset' \
                                '--baseline[Baseline experiment id]:experiment' \
                                '--candidate[Candidate experiment id]:experiment' \
                                '--cluster[Failure cluster id]:cluster' \
                                '--latest-compare[Use latest compare artifact]' \
                                '--limit[Maximum cases to replay]:limit' \
                                '--jobs[Number of cases to execute in parallel]:jobs'
                            ;;
                        *)
                            _arguments \
                                '--agent-id[Agent spec identifier]:agent_id' \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--subset[Named eval split to run]:subset' \
                                '--jobs[Number of cases to execute in parallel]:jobs' \
                                '--repeat[Run multiple times]:repeat' \
                                '--series-label[Optional label for repeated stability series]:label' \
                                '--resume-experiment[Resume a previously interrupted experiment]:experiment' \
                                '--case-id[Restrict to specific case id]:case_id' \
                                '--tag[Restrict to cases with given tag]:tag' \
                                '--match[Restrict to cases matching text]:text' \
                                '--case-limit[Maximum number of matching cases]:limit'
                            ;;
                    esac
                    ;;
                compare)
                    _arguments \
                        '--baseline[Explicit baseline experiment id]:experiment' \
                        '--best[Compare against best comparable baseline]' \
                        '--candidate[Candidate experiment id]:experiment'
                    ;;
                check)
                    _arguments \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Eval subset]:subset' \
                        '--experiment[Experiment id]:experiment' \
                        '--workflow-file[Workflow file or id]:workflow_file' \
                        '--chain-agent-id[Agent id in a chain]:agent_id' \
                        '--format[Output format]:format:(json text)' \
                        '--fail-on[Minimum severity that fails the command]:severity:(none critical high medium warning)'
                    ;;
                run-chain|run-workflow)
                    _arguments \
                        '--agent-id[Agent id]:agent_id' \
                        '--chain-agent-id[Agent id in a chain]:agent_id' \
                        '--workflow-file[Workflow file or id]:workflow_file' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset' \
                        '--case-id[Restrict to specific case id]:case_id' \
                        '--tag[Restrict to cases with given tag]:tag' \
                        '--match[Restrict to cases matching text]:text' \
                        '--limit[Maximum number of matching cases]:limit' \
                        '--case-limit[Maximum number of matching cases]:limit'
                    ;;
                status)
                    _arguments \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset' \
                        '--benchmark-signature[Benchmark signature]:benchmark_signature' \
                        '--promoted-only[Show only promoted baselines]' \
                        '--limit[Maximum number of streams]:limit'
                    ;;
                next|diagnose|propose|iterate)
                    _arguments \
                        '--experiment[Experiment id]:experiment' \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset' \
                        '--benchmark-signature[Benchmark signature]:benchmark_signature' \
                        '--baseline[Baseline experiment id]:experiment'
                    ;;
                optimize)
                    _arguments \
                        '--experiment[Experiment id]:experiment' \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset' \
                        '--benchmark-signature[Benchmark signature]:benchmark_signature' \
                        '--baseline[Baseline experiment id]:experiment' \
                        '--budget[Search budget]:budget' \
                        '--dry-run[Write a launch plan without running experiments]' \
                        '--resume[Load an existing optimization artifact]:optimize_id' \
                        '--timeout-seconds[Best-effort timeout between stages]:seconds'
                    ;;
                path-compare)
                    _arguments \
                        '--experiment[Experiment id]:experiment' \
                        '--json[Emit JSON]'
                    ;;
                adapters)
                    _arguments \
                        '--format[Output format]:format:(text json)'
                    ;;
                benchmarks)
                    _arguments \
                        '1: :->benchmarks_command'
                    case $words[3] in
                        register)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--benchmark-id[Optional benchmark identifier]:benchmark_id' \
                                '--force[Replace existing benchmark]'
                            ;;
                        list)
                            _arguments \
                                '--format[Output format]:format:(text json)'
                            ;;
                        show|check)
                            _arguments \
                                '--benchmark-id[Registered benchmark identifier]:benchmark_id' \
                                '--version[Specific benchmark version]:version'
                            ;;
                        diff)
                            _arguments \
                                '--benchmark-id[Registered benchmark identifier]:benchmark_id' \
                                '--version-a[Older version]:version' \
                                '--version-b[Newer version]:version'
                            ;;
                        install)
                            _arguments \
                                '--benchmark-id[Registered benchmark identifier]:benchmark_id' \
                                '--eval-id[Optional eval id override]:eval_id' \
                                '--force[Overwrite existing eval suite]'
                            ;;
                    esac
                    ;;
                redteam)
                    _arguments \
                        '1: :->redteam_command'
                    case $words[3] in
                        generate)
                            _arguments \
                                '--source-eval-id[Source eval suite id]:eval_id' \
                                '--output-eval-id[Output eval suite id]:eval_id' \
                                '--mode[Red-team mode]:mode'
                            ;;
                        report)
                            _arguments \
                                '--experiment[Experiment id]:experiment' \
                                '--baseline[Optional baseline experiment id]:experiment'
                            ;;
                    esac
                    ;;
                human-review)
                    _arguments \
                        '1: :->human_review_command'
                    case $words[3] in
                        queues)
                            _arguments \
                                '--pending-only[Show only queues with pending items]'
                            ;;
                        answers)
                            _arguments
                            ;;
                        answer)
                            _arguments \
                                '--queue-id[Queue identifier]:queue_id' \
                                '--checkpoint-id[Checkpoint id]:checkpoint_id' \
                                '--answer[Free-form answer]:answer' \
                                '--decision[Decision label]:decision' \
                                '--guidance[Guidance note]:guidance' \
                                '--case-id[Related case id]:case_id' \
                                '--reviewer[Reviewer label]:reviewer'
                            ;;
                        apply)
                            _arguments \
                                '--queue-id[Queue identifier]:queue_id' \
                                '--checkpoint-id[Checkpoint identifier]:checkpoint_id' \
                                '--dry-run[Show planned edits without writing]'
                            ;;
                    esac
                    ;;
                artifacts)
                    _arguments \
                        '1: :->artifact_command'
                    case $words[3] in
                        show)
                            _arguments \
                                '1: :->show_command'
                            case $words[4] in
                                experiment)
                                    _arguments \
                                        '--experiment[Experiment id]:experiment'
                                    ;;
                                case)
                                    _arguments \
                                        '--experiment[Experiment id]:experiment' \
                                        '--case-id[Case id]:case_id'
                                    ;;
                                compare)
                                    _arguments \
                                        '--baseline[Baseline experiment id]:experiment' \
                                        '--candidate[Candidate experiment id]:experiment' \
                                        '--latest[Load latest compare artifact]'
                                    ;;
                            esac
                            ;;
                        export)
                            _arguments \
                                '--experiment[Experiment id]:experiment' \
                                '--output[Optional output path]:output'
                            ;;
                        review)
                            _arguments \
                                '--experiment[Experiment id]:experiment' \
                                '--baseline[Optional baseline experiment id]:experiment' \
                                '--queue-id[Optional review queue id]:queue_id' \
                                '--output[Optional HTML output path]:output'
                            ;;
                    esac
                    ;;
                scorer)
                    _arguments \
                        '1: :->scorer_command'
                    case $words[3] in
                        template)
                            _arguments '--kind[Template kind]:kind:(heuristic dialogue)'
                            ;;
                        lint)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--path[Scorer file path]:path:_files'
                            ;;
                        schema)
                            _arguments
                            ;;
                    esac
                    ;;
                route-policy)
                    _arguments \
                        '1: :->route_policy_command'
                    case $words[3] in
                        template)
                            _arguments '--strategy[Route strategy]:strategy:(direct clarification recovery thorough)'
                            ;;
                        draft|lint|explain|benchmark)
                            _arguments \
                                '--eval-id[Eval suite identifier]:eval_id' \
                                '--agent-id[Agent id]:agent_id' \
                                '--apply[Write drafted policies]' \
                                '--force[Overwrite existing files]'
                            ;;
                        review)
                            _arguments \
                                '--experiment[Experiment id]:experiment' \
                                '--case-id[Case id]:case_id' \
                                '--decision[Review decision]:decision:(approve reject convert)' \
                                '--route-id[Route id]:route_id' \
                                '--eval-id[Eval suite id]:eval_id' \
                                '--dry-run[Show planned update]'
                            ;;
                    esac
                    ;;
                scorer-audit)
                    _arguments \
                        '--experiment[Experiment id]:experiment' \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset'
                    ;;
                release)
                    _arguments \
                        '--agent-id[Agent id]:agent_id' \
                        '--eval-id[Eval id]:eval_id' \
                        '--subset[Subset]:subset' \
                        '--benchmark-signature[Benchmark signature]:benchmark_signature' \
                        '--promoted-only[Show only promoted baselines]' \
                        '--limit[Maximum number of streams]:limit'
                    ;;
                migrate)
                    _arguments \
                        '--check[Check migrate targets without writing changes]'
                    ;;
                doctor|traces)
                    _arguments \
                        '--format[Output format]:format:(text json)'
                    ;;
            esac
            ;;
    esac
}

_awb "$@"
