# Bash completion for the `ml` CLI
# Usage:
#   source /path/to/ml_completion.bash
# or
#   echo 'source /path/to/ml_completion.bash' >> ~/.bashrc

_ml_completions()
{
    local cur prev cmds
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"

    # Global options
    global_opts="--help --verbose --quiet --monitor"

    # Top-level subcommands
    cmds="init sync queue requeue status monitor cancel prune watch dataset experiment"

    # If completing the subcommand itself
    if [[ ${COMP_CWORD} -eq 1 ]]; then
        COMPREPLY=( $(compgen -W "${cmds} ${global_opts}" -- "${cur}") )
        return 0
    fi

    # Handle global options anywhere
    case "${prev}" in
        --help|--verbose|--quiet|--monitor)
            # No further completion after global flags
            return 0
            ;;
    esac

    case "${COMP_WORDS[1]}" in
        init)
            # No specific arguments for init
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        sync)
            # Complete directories for sync
            COMPREPLY=( $(compgen -d -- "${cur}") )
            ;;
        queue)
            queue_opts="--commit --priority --cpu --memory --gpu --gpu-memory --snapshot-id --snapshot-sha256 --args -- ${global_opts}"
            case "${prev}" in
                --priority)
                    COMPREPLY=( $(compgen -W "0 1 2 3 4 5 6 7 8 9 10" -- "${cur}") )
                    ;;
                --cpu|--memory|--gpu)
                    COMPREPLY=( $(compgen -W "0 1 2 4 8 16 32" -- "${cur}") )
                    ;;
                --gpu-memory)
                    COMPREPLY=( $(compgen -W "4 8 16 24 32 48" -- "${cur}") )
                    ;;
                --commit|--snapshot-id|--snapshot-sha256|--args)
                    # Free-form; no special completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${queue_opts}" -- "${cur}") )
                    else
                        # Suggest common job names (static for now)
                        COMPREPLY=( $(compgen -W "train evaluate deploy" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        requeue)
            requeue_opts="--name --priority --cpu --memory --gpu --gpu-memory --args -- ${global_opts}"
            case "${prev}" in
                --priority)
                    COMPREPLY=( $(compgen -W "0 1 2 3 4 5 6 7 8 9 10" -- "${cur}") )
                    ;;
                --cpu|--memory|--gpu)
                    COMPREPLY=( $(compgen -W "0 1 2 4 8 16 32" -- "${cur}") )
                    ;;
                --gpu-memory)
                    COMPREPLY=( $(compgen -W "4 8 16 24 32 48" -- "${cur}") )
                    ;;
                --name|--args)
                    # Free-form; no special completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${requeue_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        status)
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        monitor)
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        cancel)
            # Job names often free-form; no special completion
            ;;
        prune)
            case "${prev}" in
                --keep)
                    COMPREPLY=( $(compgen -W "1 2 3 4 5" -- "${cur}") )
                    ;;
                --older-than)
                    COMPREPLY=( $(compgen -W "1h 6h 12h 1d 3d 7d" -- "${cur}") )
                    ;;
                *)
                    COMPREPLY=( $(compgen -W "--keep --older-than ${global_opts}" -- "${cur}") )
                    ;;
            esac
            ;;
        watch)
            # Complete directories for watch
            COMPREPLY=( $(compgen -d -- "${cur}") )
            ;;
        dataset)
            COMPREPLY=( $(compgen -W "list upload download delete info search" -- "${cur}") )
            ;;
        experiment)
            COMPREPLY=( $(compgen -W "log show" -- "${cur}") )
            ;;
        *)
            # Fallback to global options
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
    esac

    return 0
}

complete -F _ml_completions ml