# Zsh completion for the `ml` CLI
# Usage:
#   source /path/to/ml_completion.zsh
# or add to your ~/.zshrc:
#   source /path/to/ml_completion.zsh

_ml() {
  local -a subcommands
  subcommands=(
    'init:Setup configuration interactively'
    'sync:Sync project to server'
    'queue:Queue job for execution'
    'requeue:Re-submit a previous run/commit'
    'status:Get system status'
    'monitor:Launch TUI via SSH'
    'cancel:Cancel running job'
    'prune:Prune old experiments'
    'watch:Watch directory for auto-sync'
    'dataset:Manage datasets'
    'find:Search experiments by tags/outcome/dataset'
    'export:Export experiment for sharing'
    'compare:Compare two runs narrative fields'
    'outcome:Set post-run outcome'
    'info:Show run info'
    'logs:Fetch job logs'
    'annotate:Add annotation to run'
    'validate:Validate provenance'
  )

  local -a global_opts
  global_opts=(
    '--help:Show this help message'
    '--verbose:Enable verbose output'
    '--quiet:Suppress non-error output'
    '--monitor:Monitor progress of long-running operations'
    '--json:Output structured JSON'
  )

  local curcontext="$curcontext" state line
  _arguments -C \
    '1:command:->cmds' \
    '*::arg:->args'

  case $state in
    cmds)
      _describe -t commands 'ml commands' subcommands
      return
    ;;
    args)
      case $words[2] in
        sync)
          _arguments -C \
            '--help[Show sync help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]' \
            '1:directory:_directories'
        ;;
        queue)
          _arguments -C \
            '--help[Show queue help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--json[Output JSON]' \
            '--commit[Commit id (40-hex) or unique prefix (>=7)]:commit id:' \
            '--priority[Priority (0-255)]:priority:' \
            '--cpu[CPU cores]:cpu:' \
            '--memory[Memory (GB)]:memory:' \
            '--gpu[GPU count]:gpu:' \
            '--gpu-memory[GPU memory]:gpu memory:' \
            '--snapshot-id[Snapshot id]:snapshot id:' \
            '--snapshot-sha256[Snapshot sha256]:snapshot sha256:' \
            '--args[Runner args string]:args:' \
            '--note[Human notes]:note:' \
            '--hypothesis[Research hypothesis]:hypothesis:' \
            '--context[Background context]:context:' \
            '--intent[What you are trying to accomplish]:intent:' \
            '--expected-outcome[What you expect to happen]:expected outcome:' \
            '--experiment-group[Group related experiments]:experiment group:' \
            '--tags[Comma-separated tags]:tags:' \
            '--dry-run[Show what would be queued]' \
            '--validate[Validate without queuing]' \
            '--explain[Explain what will happen]' \
            '--force[Queue even if duplicate exists]' \
            '--mlflow[Enable MLflow]' \
            '--mlflow-uri[MLflow tracking URI]:uri:' \
            '--tensorboard[Enable TensorBoard]' \
            '--wandb-key[Wandb API key]:key:' \
            '--wandb-project[Wandb project]:project:' \
            '--wandb-entity[Wandb entity]:entity:' \
            '1:job name:' \
            '*:args separator:(--)'
        ;;
        requeue)
          _arguments -C \
            '--help[Show requeue help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]' \
            '--name[Override job name]:job name:' \
            '--priority[Priority (0-255)]:priority:' \
            '--cpu[CPU cores]:cpu:' \
            '--memory[Memory (GB)]:memory:' \
            '--gpu[GPU count]:gpu:' \
            '--gpu-memory[GPU memory]:gpu memory:' \
            '--args[Runner args string]:args:' \
            '1:commit_id|run_id|task_id|path:' \
            '*:args separator:(--)'
        ;;
        status)
          _arguments -C \
            '--help[Show status help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]'
        ;;
        monitor)
          _arguments -C \
            '--help[Show monitor help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]'
        ;;
        cancel)
          _arguments -C \
            '--help[Show cancel help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]' \
            '1:job name:'
        ;;
        prune)
          _arguments -C \
            '--help[Show prune help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]' \
            '--keep[Keep N most recent experiments]:number:' \
            '--older-than[Remove experiments older than D days]:days:'
        ;;
        watch)
          _arguments -C \
            '--help[Show watch help]' \
            '--verbose[Enable verbose output]' \
            '--quiet[Suppress non-error output]' \
            '--monitor[Monitor progress]' \
            '1:directory:_directories'
        ;;
        dataset)
          _values 'dataset subcommand' \
            'list[list datasets]' \
            'register[register dataset]' \
            'info[dataset info]' \
            'search[search datasets]' \
            'verify[verify dataset]'
        ;;
        experiment)
          _values 'experiment subcommand' \
            'log[log metrics]' \
            'show[show experiment]' \
            'list[list experiments]'
        ;;
        narrative)
          _arguments -C \
            '1:subcommand:(set)' \
            '--hypothesis[Research hypothesis]:hypothesis:' \
            '--context[Background context]:context:' \
            '--intent[What you are trying to accomplish]:intent:' \
            '--expected-outcome[What you expect to happen]:expected outcome:' \
            '--parent-run[Parent run ID]:parent run:' \
            '--experiment-group[Group related experiments]:experiment group:' \
            '--tags[Comma-separated tags]:tags:' \
            '--base[Base path]:base:_directories' \
            '--json[Output JSON]' \
            '--help[Show help]'
        ;;
        outcome)
          _arguments -C \
            '1:subcommand:(set)' \
            '--outcome[Outcome status]:outcome:(validates refutes inconclusive partial)' \
            '--summary[Summary of results]:summary:' \
            '--learning[A learning from this run]:learning:' \
            '--next-step[Suggested next step]:next step:' \
            '--validation-status[Did results validate hypothesis]:(validates refutes inconclusive)' \
            '--surprise[Unexpected finding]:surprise:' \
            '--base[Base path]:base:_directories' \
            '--json[Output JSON]' \
            '--help[Show help]'
        ;;
        info|logs|annotate|validate)
          _arguments -C \
            '--base[Base path]:base:_directories' \
            '--json[Output JSON]' \
            '--help[Show help]' \
            '1:run id or path:'
        ;;
        compare)
          _arguments -C \
            '--help[Show compare help]' \
            '--json[Output JSON]' \
            '--csv[Output CSV for spreadsheets]' \
            '--all[Show all fields]' \
            '--fields[Specify fields to compare]:fields:(narrative metrics metadata outcome)' \
            '1:run a:' \
            '2:run b:'
        ;;
        find)
          _arguments -C \
            '--help[Show find help]' \
            '--json[Output JSON]' \
            '--csv[Output CSV for spreadsheets]' \
            '--limit[Max results]:limit:(10 20 50 100)' \
            '--tag[Filter by tag]:tag:' \
            '--outcome[Filter by outcome]:outcome:(validates refutes inconclusive partial)' \
            '--dataset[Filter by dataset]:dataset:' \
            '--experiment-group[Filter by group]:group:' \
            '--author[Filter by author]:author:' \
            '--after[After date]:date:' \
            '--before[Before date]:date:' \
            '::query:'
        ;;
        export)
          _arguments -C \
            '--help[Show export help]' \
            '--json[Output JSON]' \
            '--bundle[Create bundle at path]:path:_files' \
            '--anonymize[Enable anonymization]' \
            '--anonymize-level[Anonymization level]:level:(metadata-only full)' \
            '--base[Base path]:base:_directories' \
            '1:run id or path:'
        ;;
        *)
          _arguments -C "${global_opts[@]}"
        ;;
      esac
    ;;
  esac
}

compdef _ml ml