#!/usr/bin/env bash

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

    # Subcommands
    subcmds="virtual mig"

    # If cache variable _HY_SMI_OPTS is empty, parse options from help output.
    if [[ -z "$_HY_SMI_OPTS" ]]; then
        # Extract only valid short/long options (e.g. -h, --help),
        # avoiding punctuation-contaminated tokens like -h' or --foo.
        export _HY_SMI_OPTS=$(hy-smi -h 2>/dev/null | awk '{
            for(i=1;i<=NF;i++) {
                token=$i
                while (match(token, /--[A-Za-z0-9][A-Za-z0-9-]*|-[A-Za-z0-9]/)) {
                    print substr(token, RSTART, RLENGTH)
                    token=substr(token, RSTART + RLENGTH)
                }
            }
        }' | sort -u)
    fi

    # Provide intelligent suggestions for specific parameters (prev)
    case "${prev}" in
        --load|--save|--filename|-fdmon)
            COMPREPLY=( $(compgen -f -- "${cur}") )
            return 0
            ;;
        --loglevel)
            COMPREPLY=( $(compgen -W "debug info warning error critical" -- "${cur}") )
            return 0
            ;;
        --setperflevel)
            COMPREPLY=( $(compgen -W "auto low high manual" -- "${cur}") )
            return 0
            ;;
        --setecc)
            COMPREPLY=( $(compgen -W "on off stable" -- "${cur}") )
            return 0
            ;;
        --direction)
            COMPREPLY=( $(compgen -W "0 1" -- "${cur}") )
            return 0
            ;;
        --enablelowpower|--setautoloaddriver|--setautorunhyckptd|-mig|--multi-instance-gpu)
            COMPREPLY=( $(compgen -W "0 1" -- "${cur}") )
            return 0
            ;;
        --setboost)
            COMPREPLY=( $(compgen -W "0-AI(default), 1-HPC(default), 2-AI(typical), 3-HPC(typical)" -- "${cur}") )
            return 0
            ;;
    esac

    if [[ ${cur} == -* ]] ; then
        # Only complete the dynamically extracted flags
        COMPREPLY=( $(compgen -W "${_HY_SMI_OPTS}" -- "${cur}") )
        return 0
    fi

    # Default complete all flags and subcommands
    COMPREPLY=( $(compgen -W "${_HY_SMI_OPTS} ${subcmds}" -- "${cur}") )
}

complete -o default -F _hy_smi_completion hy-smi