#!/usr/bin/env bash

# Bash completion for slurm-quota

_slurm_quota_list_accounts() {
    if command -v sacctmgr >/dev/null 2>&1; then
        sacctmgr -nP show account format=account 2>/dev/null \
            | cut -d'|' -f1 \
            | sort -u
    fi
}

_slurm_quota_complete_users() {
    compgen -u -- "$1"
}

_slurm_quota_complete_accounts() {
    local cur="$1"
    local accounts
    accounts="$(_slurm_quota_list_accounts)"
    if [[ -n "$accounts" ]]; then
        compgen -W "$accounts" -- "$cur"
    fi
}

_slurm_quota() {
    local cur prev cword
    local i
    local -a words

    cur="${COMP_WORDS[COMP_CWORD]}"
    prev=""
    if (( COMP_CWORD > 0 )); then
        prev="${COMP_WORDS[COMP_CWORD-1]}"
    fi
    cword="$COMP_CWORD"
    words=("${COMP_WORDS[@]}")

    local -a subcommands
    subcommands=(
        login
        token
        stats
        role
        adjust
        user-quota
        account-quota
        user-gpu-quota
        account-gpu-quota
        gpu-factors
        set-gpu-factor
        default-quotas
        set-default-quotas
    )

    local cmd_idx=0
    for (( i=0; i<=cword; i++ )); do
        if [[ "${words[i]}" == "slurm-quota" || "${words[i]}" == */slurm-quota ]]; then
            cmd_idx="$i"
            break
        fi
    done

    local subcmd=""
    local subcmd_idx=-1
    # Parse only words after command name and before the one being completed.
    # Including the current partial token here prevents proper autocompletion.
    for (( i=cmd_idx+1; i<cword; i++ )); do
        if [[ "${words[i]}" == --debug || "${words[i]}" == -q || "${words[i]}" == --quiet || "${words[i]}" == --version || "${words[i]}" == -h || "${words[i]}" == --help ]]; then
            continue
        fi
        if [[ "${words[i]}" == --log-flags || "${words[i]}" == --debug-flags ]]; then
            ((i++))
            continue
        fi
        if [[ "${words[i]}" != -* ]]; then
            subcmd="${words[i]}"
            subcmd_idx="$i"
            break
        fi
    done

    if [[ -z "$subcmd" ]]; then
        COMPREPLY=($(compgen -W "--debug -q --quiet --log-flags --debug-flags --version -h --help ${subcommands[*]}" -- "$cur"))
        return 0
    fi

    case "$prev" in
        --user)
            COMPREPLY=($(_slurm_quota_complete_users "$cur"))
            return 0
            ;;
        --account)
            COMPREPLY=($(_slurm_quota_complete_accounts "$cur"))
            return 0
            ;;
        --minutes|--hours|--user-cpu|--user-gpu|--account-cpu|--account-gpu|--reason)
            return 0
            ;;
    esac

    local base_opts="-h --help --version"
    case "$subcmd" in
        gpu-factors|default-quotas)
            COMPREPLY=($(compgen -W "$base_opts" -- "$cur"))
            ;;
        login)
            if (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(_slurm_quota_complete_users "$cur"))
            else
                COMPREPLY=($(compgen -W "$base_opts --save" -- "$cur"))
            fi
            ;;
        token)
            COMPREPLY=($(compgen -W "$base_opts --save" -- "$cur"))
            ;;
        stats)
            COMPREPLY=($(compgen -W "$base_opts --user --account --all --hours" -- "$cur"))
            COMPREPLY+=($(_slurm_quota_complete_users "$cur"))
            ;;
        adjust)
            COMPREPLY=($(compgen -W "$base_opts --user --account --cpu --gpu --minutes --hours --reason" -- "$cur"))
            ;;
        user-quota|user-gpu-quota)
            if [[ "$cur" == -* ]]; then
                COMPREPLY=($(compgen -W "$base_opts --reason" -- "$cur"))
            elif (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(_slurm_quota_complete_users "$cur"))
            elif (( cword == subcmd_idx + 2 )); then
                COMPREPLY=($(compgen -W "-1" -- "$cur"))
            else
                COMPREPLY=($(compgen -W "$base_opts --reason" -- "$cur"))
            fi
            ;;
        account-quota|account-gpu-quota)
            if [[ "$cur" == -* ]]; then
                COMPREPLY=($(compgen -W "$base_opts --reason" -- "$cur"))
            elif (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(_slurm_quota_complete_accounts "$cur"))
            elif (( cword == subcmd_idx + 2 )); then
                COMPREPLY=($(compgen -W "-1" -- "$cur"))
            else
                COMPREPLY=($(compgen -W "$base_opts --reason" -- "$cur"))
            fi
            ;;
        set-gpu-factor)
            if (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(compgen -W "default" -- "$cur"))
            fi
            ;;
        set-default-quotas)
            COMPREPLY=($(compgen -W "$base_opts --user-cpu --user-gpu --account-cpu --account-gpu" -- "$cur"))
            ;;
        role)
            local role_idx=-1 role_subcmd="" role_subcmd_idx=-1
            local managers_subcmd="" managers_subcmd_idx=-1
            for (( i=cmd_idx+1; i<cword; i++ )); do
                if [[ "${words[i]}" == --debug || "${words[i]}" == -q || "${words[i]}" == --quiet || "${words[i]}" == --version || "${words[i]}" == -h || "${words[i]}" == --help ]]; then
                    continue
                fi
                if [[ "${words[i]}" == --log-flags || "${words[i]}" == --debug-flags ]]; then
                    ((i++))
                    continue
                fi
                if [[ "${words[i]}" == -* ]]; then
                    continue
                fi
                if (( role_idx < 0 )); then
                    if [[ "${words[i]}" == "role" ]]; then
                        role_idx="$i"
                    fi
                    continue
                fi
                if [[ -z "$role_subcmd" ]]; then
                    role_subcmd="${words[i]}"
                    role_subcmd_idx="$i"
                    continue
                fi
                if [[ "$role_subcmd" == "managers" && -z "$managers_subcmd" ]]; then
                    managers_subcmd="${words[i]}"
                    managers_subcmd_idx="$i"
                fi
            done

            if (( cword == role_idx + 1 )); then
                COMPREPLY=($(compgen -W "$base_opts show list grant revoke managers" -- "$cur"))
            elif [[ "$role_subcmd" == "show" || "$role_subcmd" == "list" ]]; then
                COMPREPLY=($(compgen -W "$base_opts" -- "$cur"))
            elif [[ "$role_subcmd" == "grant" || "$role_subcmd" == "revoke" ]]; then
                if (( cword == role_subcmd_idx + 1 )); then
                    COMPREPLY=($(compgen -W "operator manager" -- "$cur"))
                elif (( cword == role_subcmd_idx + 2 )); then
                    COMPREPLY=($(_slurm_quota_complete_users "$cur"))
                else
                    COMPREPLY=($(compgen -W "$base_opts" -- "$cur"))
                fi
            elif [[ "$role_subcmd" == "managers" ]]; then
                if (( cword == role_subcmd_idx + 1 )); then
                    COMPREPLY=($(compgen -W "list add remove" -- "$cur"))
                elif [[ "$managers_subcmd" == "list" ]]; then
                    if (( cword == managers_subcmd_idx + 1 )); then
                        COMPREPLY=($(_slurm_quota_complete_users "$cur"))
                    else
                        COMPREPLY=($(compgen -W "$base_opts" -- "$cur"))
                    fi
                elif [[ "$managers_subcmd" == "add" || "$managers_subcmd" == "remove" ]]; then
                    if (( cword == managers_subcmd_idx + 1 )); then
                        COMPREPLY=($(_slurm_quota_complete_users "$cur"))
                    elif (( cword == managers_subcmd_idx + 2 )); then
                        COMPREPLY=($(_slurm_quota_complete_accounts "$cur"))
                    else
                        COMPREPLY=($(compgen -W "$base_opts" -- "$cur"))
                    fi
                else
                    COMPREPLY=($(compgen -W "$base_opts list add remove" -- "$cur"))
                fi
            else
                COMPREPLY=($(compgen -W "$base_opts show list grant revoke managers" -- "$cur"))
            fi
            ;;
    esac

    return 0
}

complete -F _slurm_quota slurm-quota
