#!/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 j
    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 public_subcommands restricted_subcommands subcommands
    public_subcommands=(
        stats
        gpu-factors
        default-quotas
    )
    restricted_subcommands=(
        adjust
        user-quota
        account-quota
        user-gpu-quota
        account-gpu-quota
        set-gpu-factor
        set-default-quotas
        prune
    )
    subcommands=("${public_subcommands[@]}")

    local is_privileged=0
    if [[ "$EUID" -eq 0 || "$USER" == "slurm" ]]; then
        is_privileged=1
    fi

    # Allow privileged command completion only for literal "sudo" in first position.
    if [[ "${words[0]}" == "sudo" ]]; then
        is_privileged=1
    fi

    if (( is_privileged )); then
        subcommands+=("${restricted_subcommands[@]}")
    fi

    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]}" == "sudo" ]]; then
            continue
        fi
        if [[ "${words[i]}" == "-u" || "${words[i]}" == "-g" ]]; then
            (( i++ ))
            continue
        fi
        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]}" != -* ]]; then
            subcmd="${words[i]}"
            subcmd_idx="$i"
            break
        fi
    done

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

    # Intentionally do not autocomplete internal/service subcommands.
    if [[ "$subcmd" == "charge" || "$subcmd" == "serve" ]]; then
        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|--port|--idle-timeout|--host|--user-cpu|--user-gpu|--account-cpu|--account-gpu)
            return 0
            ;;
    esac

    local base_opts="-h --help --version"
    case "$subcmd" in
        charge|gpu-factors|default-quotas)
            COMPREPLY=($(compgen -W "$base_opts" -- "$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" -- "$cur"))
            ;;
        user-quota|user-gpu-quota)
            if (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(_slurm_quota_complete_users "$cur"))
            elif (( cword == subcmd_idx + 2 )); then
                COMPREPLY=($(compgen -W "-1" -- "$cur"))
            fi
            ;;
        account-quota|account-gpu-quota)
            if (( cword == subcmd_idx + 1 )); then
                COMPREPLY=($(_slurm_quota_complete_accounts "$cur"))
            elif (( cword == subcmd_idx + 2 )); then
                COMPREPLY=($(compgen -W "-1" -- "$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"))
            ;;
        prune)
            COMPREPLY=($(compgen -W "$base_opts --preallocs --users --accounts --all --dry-run --user --account" -- "$cur"))
            ;;
        serve)
            COMPREPLY=($(compgen -W "$base_opts --host --port --idle-timeout" -- "$cur"))
            ;;
    esac

    return 0
}

complete -F _slurm_quota slurm-quota
