# 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 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)
            # Suggest common job names (static for now)
            COMPREPLY=( $(compgen -W "train evaluate deploy" -- "${cur}") )
            ;;
        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