#compdef nats

# Zsh completion for nats
# Dynamically generated from --completion-zsh-menu output
#
# Installation:
#   Option 1: Copy to your fpath
#     nats --completion-script-zsh > /usr/local/share/zsh/site-functions/_nats
#
#   Option 2: Source in .zshrc
#     source <(nats --completion-script-zsh)
#
#   Option 3: Save to a file and add its directory to fpath in .zshrc (before compinit)
#     fpath=(/path/to/completions $fpath)
#     autoload -Uz compinit && compinit

# Enable menu selection with highlighting
# Preserve the order emitted by the command (most common first)
zmodload zsh/complist
zstyle ':completion:*:nats:*' menu select=1 interactive
zstyle ':completion:*:nats:*' sort false
zstyle ':completion:*:nats:*' group-name ''
zstyle ':completion:*:nats:*:descriptions' format '%B%F{cyan}-- %d --%f%b'

# Matching: case-insensitive, then substring, then partial-word at hyphens
zstyle ':completion:*:nats:*' matcher-list \
    'm:{a-zA-Z}={A-Za-z}' \
    'm:{a-zA-Z}={A-Za-z} l:|=* r:|=*' \
    'm:{a-zA-Z}={A-Za-z} r:|[-]=** l:|=*'

_nats() {
    local curcontext="$curcontext"
    local cur prev flagname
    local -i i j skip_next pos_count arg_idx cmd_depth
    local lt field1 field2 field3 field4 field5 field6
    local fname fshort fplace fneg fdesc fenums

    local -a cmd_path
    local -a cmds cdesc
    local -a fnames fshorts fplaceholders fnegatables fdescs fenumvals
    local -a anames arequireds adescs
    local -a flag_list cmd_list enum_list

    # Build the command path: extract non-flag words between the program
    # name (word 1) and the cursor position.  Flags and their values are
    # skipped so that we resolve the right subcommand context.
    i=2 skip_next=0
    while (( i < CURRENT )); do
        if (( skip_next )); then
            skip_next=0
            (( i++ ))
            continue
        fi
        case "${words[i]}" in
            --*=*) ;;           # --flag=value is self-contained
            --*)                # long flag, maybe followed by a value
                if (( i + 1 < CURRENT )) && [[ "${words[i+1]}" != -* ]]; then
                    skip_next=1
                fi ;;
            -?)                 # single short flag like -s
                if (( i + 1 < CURRENT )) && [[ "${words[i+1]}" != -* ]]; then
                    skip_next=1
                fi ;;
            *)  cmd_path+=("${words[i]}") ;;
        esac
        (( i++ ))
    done

    # Fetch machine-readable completion data for the resolved context
    local comp_output
    comp_output="$(nats "${cmd_path[@]}" --completion-zsh-menu 2>/dev/null)"
    [[ -z "$comp_output" ]] && return 1

    # Parse the tab-separated output into arrays
    cmd_depth=0
    while IFS=$'\t' read -r lt field1 field2 field3 field4 field5 field6; do
        case "$lt" in
            D) cmd_depth=$field1 ;;
            C) cmds+=("$field1");   cdesc+=("$field2") ;;
            F) fnames+=("$field1"); fshorts+=("$field2")
               fplaceholders+=("$field3"); fnegatables+=("$field4")
               fdescs+=("$field5"); fenumvals+=("$field6") ;;
            A) anames+=("$field1"); arequireds+=("$field2")
               adescs+=("$field3") ;;
        esac
    done <<< "$comp_output"

    cur="${words[CURRENT]}"

    # --- Check if we are completing a flag value ---
    # If the previous word is a flag that takes a value, complete the value.
    if (( CURRENT > 2 )); then
        prev="${words[CURRENT-1]}"
        if [[ "$prev" == --* || "$prev" == -? ]]; then
            flagname=""
            if [[ "$prev" == --* ]]; then
                flagname="${prev#--}"
            else
                # Find the long name for this short flag
                for (( j = 1; j <= ${#fnames}; j++ )); do
                    if [[ "-${fshorts[j]}" == "$prev" ]]; then
                        flagname="${fnames[j]}"
                        break
                    fi
                done
            fi
            if [[ -n "$flagname" ]]; then
                for (( j = 1; j <= ${#fnames}; j++ )); do
                    if [[ "${fnames[j]}" == "$flagname" && "${fplaceholders[j]}" != "_" ]]; then
                        # If this flag has enum values, offer them as completions
                        if [[ -n "${fenumvals[j]}" ]]; then
                            enum_list=()
                            for fenums in ${(s:,:)fenumvals[j]}; do
                                enum_list+=("$fenums")
                            done
                            _describe -t enum-values "${fplaceholders[j]}" enum_list -o nosort && return 0
                        fi
                        case "${fplaceholders[j]}" in
                            FILE) _files && return 0 ;;
                            DIR)  _directories && return 0 ;;
                            *)    _message -r "${fplaceholders[j]}" && return 0 ;;
                        esac
                    fi
                done
            fi
        fi
    fi

    # --- Build the flag completion list (reused below) ---
    for (( j = 1; j <= ${#fnames}; j++ )); do
        fname="${fnames[j]}"
        fshort="${fshorts[j]}"
        fneg="${fnegatables[j]}"
        fdesc="${fdescs[j]}"

        [[ "$fname" == "help" || "$fname" == "version" ]] && continue
        [[ "$fname" == completion-* ]] && continue

        flag_list+=("--${fname}:${fdesc}")

        if [[ "$fneg" == "true" ]]; then
            flag_list+=("--no-${fname}:Disable ${fname}")
        fi

        if [[ "$fshort" != "_" ]]; then
            flag_list+=("-${fshort}:${fdesc}")
        fi
    done

    # --- Flag completion (current word starts with -) ---
    if [[ "$cur" == -* ]]; then
        _describe -t flags 'flags' flag_list -o nosort && return 0
        return 1
    fi

    # --- Command / subcommand completion ---
    if (( ${#cmds} > 0 )); then
        for (( j = 1; j <= ${#cmds}; j++ )); do
            [[ "${cmds[j]}" == "help" ]] && continue
            cmd_list+=("${cmds[j]}:${cdesc[j]}")
        done

        _describe -t commands 'commands' cmd_list -o nosort && return 0
    fi

    # --- Argument completion ---
    # Use command depth from the D line to know how many non-flag words
    # are commands vs positional args.  cmd_path contains ALL non-flag
    # words (commands + args), but only cmd_depth of them are commands.
    pos_count=$(( ${#cmd_path} - cmd_depth ))

    arg_idx=$(( pos_count + 1 ))

    if (( arg_idx >= 1 && arg_idx <= ${#anames} )); then
        _message -r "${anames[arg_idx]}: ${adescs[arg_idx]}"
        return 0
    fi

    return 1
}

if [[ "$(basename -- ${(%):-%x})" != "_nats" ]]; then
    compdef _nats nats
fi
