#!/usr/bin/env bash

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

    # Subcommands (distinct from --mig/--multi-instance-gpu set option)
    subcmds="virtual mig"

    # Lazily parse options from help output on first invocation.
    # To force a refresh after upgrading hy-smi, run: unset _HY_SMI_OPTS
    if [[ -z "$_HY_SMI_OPTS" ]]; then
        # Match both GNU-style long options (--foo-bar) and
        # tool-specific single-dash multi-char options (-idmon, -sdmon, etc.)
        _HY_SMI_OPTS=$(hy-smi -h 2>/dev/null \
            | grep -oP '(?<!\w)(--[A-Za-z][A-Za-z0-9-]*|-[A-Za-z][A-Za-z0-9]*)' \
            | sort -u)
    fi

    # Provide completions for options that take specific values
    case "${prev}" in
        # File arguments
        --load|--save|--filename|-fdmon)
            COMPREPLY=( $(compgen -f -- "${cur}") )
            return 0
            ;;

        # Enumerated string values
        --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
            ;;
        -tdmon)
            COMPREPLY=( $(compgen -W "text csv" -- "${cur}") )
            return 0
            ;;
        -sdmon)
            # Flags can be combined (e.g. cfpu), offer individual letters
            COMPREPLY=( $(compgen -W "c f p u" -- "${cur}") )
            return 0
            ;;

        # Binary 0/1 toggles
        --direction)
            COMPREPLY=( $(compgen -W "0 1" -- "${cur}") )
            return 0
            ;;
        --enablelowpower|--setautoloaddriver|--setautorunhyckptd|-mig|--multi-instance-gpu)
            COMPREPLY=( $(compgen -W "0 1" -- "${cur}") )
            return 0
            ;;

        # Multi-value enumerations
        --setboost)
            # Valid levels per help output: 0-AI(default), 2-AI(typical), 3-HPC(typical)
            # NOTE: 1 does NOT exist
            COMPREPLY=( $(compgen -W "0 2 3" -- "${cur}") )
            return 0
            ;;
        --setextendbw)
            # 0: off, 1: 1GB, 2: 2GB, 3: 4GB
            COMPREPLY=( $(compgen -W "0 1 2 3" -- "${cur}") )
            return 0
            ;;
        --channel)
            # 0-31 or all
            COMPREPLY=( $(compgen -W "$(seq 0 31) all" -- "${cur}") )
            return 0
            ;;
        --link)
            # 0-6 or all
            COMPREPLY=( $(compgen -W "0 1 2 3 4 5 6 all" -- "${cur}") )
            return 0
            ;;
    esac

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

complete -o default -F _hy_smi_completion hy-smi