# 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'
    'experiment:Manage experiments'
  )

  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'
  )

  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]' \
            '--monitor[Monitor progress]' \
            '--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:' \
            '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]' \
            'upload[upload dataset]' \
            'download[download dataset]' \
            'delete[delete dataset]' \
            'info[dataset info]' \
            'search[search datasets]'
        ;;
        experiment)
          _values 'experiment subcommand' \
            'log[log metrics]' \
            'show[show experiment]'
        ;;
        *)
          _arguments -C "${global_opts[@]}"
        ;;
      esac
    ;;
  esac
}

compdef _ml ml