#!/usr/bin/env bash

# Anthon Open Source Community 
# Pinyin Completion Hook for Bash-Completion

# Detect bash-completion
if ! declare -F _comp_compgen__call_builtin &>/dev/null; then
    echo "No function _comp_compgen__call_builtin found. Please install bash-completion first."
    exit 1
fi

# Backup the original function
eval "function __bak_comp_compgen__call_builtin() { $(declare -f _comp_compgen__call_builtin | tail -n +2) }"

# When Bash falls back to the minimal default completion handler (typically
# triggered for commands invoked via path prefixes such as ./foo or ../bar),
# bash-completion stops before calling compgen and our pinyin hook never runs.
# Wrap that fallback so path-based commands still go through _filedir and, in
# turn, our enhanced matcher.
if declare -F _comp_complete_minimal &>/dev/null && \
   ! declare -F __bak_comp_complete_minimal &>/dev/null; then
    eval "function __bak_comp_complete_minimal() { $(declare -f _comp_complete_minimal | tail -n +2) }"

    _comp_complete_minimal() {
        if [[ ${COMP_WORDS[0]+_} ]] && [[ ${COMP_WORDS[0]} == */* ]]; then
            local cur prev words cword comp_args
            _comp_initialize -- "$@" || return
            _filedir
            return
        fi

        __bak_comp_complete_minimal "$@"
    }
fi

# Decide whether a generator invoked via _comp_xfunc appears to provide
# local filesystem candidates.
_pinyin_is_local_file_generator() {
    local generator="$1"
    [[ -z "$generator" ]] && return 1

    case "$generator" in
        compgen_local_files|compgen_local_dirs|compgen_local_filedirs|compgen_filedir|compgen_filedirs)
            return 0 ;;
        compgen_local_*|*_local_files|*_local_dirs)
            return 0 ;;
    esac

    [[ "$generator" == *local*file* ]] && return 0
    [[ "$generator" == *file*local* ]] && return 0

    return 1
}

# Hook _comp_xfunc so any command-level generator that looks like it enumerates
# local files automatically benefits from the pinyin matcher.
if declare -F _comp_xfunc &>/dev/null && \
   ! declare -F __bak_comp_xfunc &>/dev/null; then
    eval "function __bak_comp_xfunc() { $(declare -f _comp_xfunc | tail -n +2) }"

    _comp_xfunc() {
        local -a __pinyin_args=("$@")
        __bak_comp_xfunc "${__pinyin_args[@]}"
        local status=$?

        local generator="${__pinyin_args[1]-}"
        if _pinyin_is_local_file_generator "$generator"; then
            local -a dir_flag=()
            local arg
            for arg in "${__pinyin_args[@]:2}"; do
                if [[ "$arg" == "-d" ]]; then
                    dir_flag=('-d')
                    break
                fi
            done
            _add_completion "${dir_flag[@]}"
        fi

        return $status
    }
fi

# Expand environment references ("$VAR", "${VAR}") inside completion prefixes
__expand_env_refs() {
    local input="$1"
    # accumulate the expanded characters here
    local result=""
    local len=${#input}
    local i=0

    # Walk through every character
    while (( i < len )); do
        local ch="${input:i:1}"
        # Preserve escaped characters verbatim
        # "\$HOME" stays "$HOME"
        if [[ "$ch" == "\\" ]]; then
            ((i++))
            if (( i < len )); then
                result+="${input:i:1}"
                ((i++))
            fi
            continue
        fi

        # Handle environment references beginning with '$'
        if [[ "$ch" == '$' ]]; then
            ((i++))
            if (( i >= len )); then
                result+='$'
                break
            fi

            ch="${input:i:1}"
            # Expand braces style 
            # "${VAR}"
            if [[ "$ch" == '{' ]]; then
                ((i++))
                local start=$i
                # Consume [A-Za-z0-9_] until we hit '}' or something else
                while (( i < len )); do
                    ch="${input:i:1}"
                    if [[ "$ch" =~ [A-Za-z0-9_] ]]; then
                        ((i++))
                        continue
                    fi
                    break
                done
                local var_name="${input:start:i-start}"
                if [[ -z "$var_name" ]]; then
                    result+='$'
                    result+='{'
                    continue
                fi
                if (( i < len )) && [[ "${input:i:1}" == '}' ]]; then
                    ((i++))
                    # ${VAR} -> ${!VAR-} returns value or empty string.
                    result+="${!var_name-}"
                else
                    result+='$'
                    result+='{'
                    i=$start
                fi
                continue
            fi

            if [[ "$ch" =~ [A-Za-z_] ]]; then
                local start=$i
                while (( i < len )) && [[ "${input:i:1}" =~ [A-Za-z0-9_] ]]; do
                    ((i++))
                done
                local var_name="${input:start:i-start}"
                result+="${!var_name-}"
                continue
            fi

            # Handle positional parameters ($1) 
            # or a few common special vars
            if [[ "$ch" =~ [0-9@*#?] ]]; then
                local special_name="$ch"
                ((i++))
                result+="${!special_name-}"
                continue
            fi

            # Any other symbol
            result+='$'
            continue
        fi

        # Normal characters are copied through.
        result+="$ch"
        ((i++))
    done

    printf '%s' "$result"
}


_comp_compgen__call_builtin() {
    __bak_comp_compgen__call_builtin "$@"
    local original_result=$?
    
    # Only add pinyin completion for file/directory completions
    local is_file_completion=false
    local compgen_args=("$@")
    
    # Check for file completion indicators
    local idx=0
    while [[ $idx -lt ${#compgen_args[@]} ]]; do
        local arg="${compgen_args[$idx]}"
        case "$arg" in
            -f|-d|-Afile|-Adirectory)
                is_file_completion=true
                break
                ;;
            -A)
                if (( idx + 1 < ${#compgen_args[@]} )); then
                    local next_arg="${compgen_args[$((idx + 1))]}"
                    if [[ "$next_arg" == "file" || "$next_arg" == "directory" ]]; then
                        is_file_completion=true
                        break
                    fi
                fi
                ;;
            file|directory)
                is_file_completion=true
                break
                ;;
        esac
        ((idx++))
    done
    
    # Also check if -W option is used with ${files[@]} or similar array expansion
    # which is a common pattern for file completion
    if [[ "$is_file_completion" == false ]]; then
        local i=0
        while [[ $i -lt ${#compgen_args[@]} ]]; do
            if [[ "${compgen_args[$i]}" == "-W" ]]; then
                local next_idx=$((i + 1))
                if [[ $next_idx -lt ${#compgen_args[@]} ]]; then
                    local word_arg="${compgen_args[$next_idx]}"
                    # Check if it contains ${files or similar array patterns
                    if [[ "$word_arg" == *'${files'* ]] || [[ "$word_arg" == *'${toks'* ]] || [[ "$word_arg" == *'$files'* ]] || [[ "$word_arg" == *'$toks'* ]]; then
                        is_file_completion=true
                        break
                    fi
                fi
            fi
            ((i++))
        done
    fi
    
    # If this looks like file completion, add pinyin matches
    if [[ "$is_file_completion" == true ]]; then
        _add_completion "$@"
    fi
    
    return $original_result
}

# Function to add completion results
_add_completion() {
    # cur: bash-completion's working value for the current word.
    local cur

    if [[ -n ${_cur-} ]]; then
        eval "cur=${_cur}"
    elif [[ -n ${COMP_WORDS+x} && -n ${COMP_CWORD+x} ]]; then
        cur="${COMP_WORDS[COMP_CWORD]-}"
    else
        return
    fi
    # origin_cur: the user's raw buffer text before any expansion
    # including quotes or ~user prefixes. in other word, "snapshot".
    local orig_cur="$cur"
    # stripped_orig: an editable copy of orig_cur
    # used to compute orig_dirpart without mutating the original text.
    local stripped_orig="$orig_cur"
    local orig_dirpart=""
    if [[ "$stripped_orig" == "'"* || "$stripped_orig" == '"'* ]]; then
        stripped_orig="${stripped_orig:1}"
    fi
    if [[ "$stripped_orig" == */* ]]; then
        orig_dirpart="${stripped_orig%/*}"
    fi
    if [[ "$orig_dirpart" == "." && "${stripped_orig:0:2}" != "./" ]]; then
        orig_dirpart=""
    fi
    
    # Check if we have the necessary variables
    local var_name
    if [[ -n ${_var-} ]]; then
        var_name="$_var"
    else
        var_name="COMPREPLY"
    fi
    
    # Skip empty
    [[ -z "$cur" ]] && return

    # perform bash-completion's normal expansions.
    _expand || return 0

    local dirpart basepart
    if [[ "${cur:0:1}" == "'" || "${cur:0:1}" == "\"" ]]; then
        dirpart="$(dirname -- "${cur:1}")"
        basepart="$(basename -- "${cur:1}")"
    else
        dirpart="$(dirname -- "$cur")"
        basepart="$(basename -- "$cur")"
    fi
    
    [[ "$dirpart" == "." && "${cur:0:2}" != "./" ]] && dirpart=""

    # Expand environemnt variables
    # dirpart_lookup: save the true path after expanded
    # NOTE: in the end, the path prefix will be rollbacked to "snapshot".
    local dirpart_lookup="$dirpart"
    if [[ -n "$dirpart_lookup" && "$dirpart_lookup" == *'$'* ]]; then
        local expanded_lookup
        expanded_lookup="$(__expand_env_refs "$dirpart_lookup")"
        if [[ -n "$expanded_lookup" ]]; then
            dirpart_lookup="$expanded_lookup"
        fi
    fi
    
    local savedPWD="$PWD"
    local resolved_dir
    local compgen_opts=(-f)
    
    local is_dir_only=false
    for arg in "$@"; do
        if [[ "$arg" == "-d" ]]; then
            is_dir_only=true
            compgen_opts=(-d)
            break
        fi
    done
    
    if [[ -n "$dirpart_lookup" ]]; then
        # Resolve the working directory for compgen use realpath, but remember
        # the original textual prefix so completions can stay aligned with what the user typed.
        resolved_dir="$(realpath -- "$dirpart_lookup" 2>/dev/null)"
        if [[ -d "$resolved_dir" ]]; then
            cd -- "$resolved_dir" 2>/dev/null || return
        else
            cd "$savedPWD" || return
            return
        fi
    fi

    # Kernel
    local -a pinyin_matched
    if [[ "$is_dir_only" == true ]]; then
        mapfile -t pinyin_matched < <(
            compgen -d -- 2>/dev/null |
            bash-pinyin-completion-rs "$basepart" 2>/dev/null
        )
    else
        mapfile -t pinyin_matched < <(
            compgen -f -- 2>/dev/null |
            bash-pinyin-completion-rs "$basepart" 2>/dev/null
        )
    fi
    
    # Restore directory
    cd "$savedPWD" || return
    
    if [[ ${#pinyin_matched[@]} -gt 0 ]]; then
        local display_dirpart="$dirpart"
        if [[ -n "$orig_dirpart" ]]; then
            # When the user typed something like ~user/src, prefer their original prefix for display
            # instead of the realpath directory we temp into.
            # "snapshot" we saved before comes in handy here.
            display_dirpart="$orig_dirpart"
        fi
        if [[ -n "$display_dirpart" ]]; then
            local sep="/"
            [[ "$display_dirpart" == "/" ]] && sep=""
            for i in "${!pinyin_matched[@]}"; do
                pinyin_matched[$i]="${display_dirpart}${sep}${pinyin_matched[$i]}"
            done
        fi

        local orig_check="$orig_cur"
        if [[ "$orig_check" == "'"* || "$orig_check" == '"'* ]]; then
            orig_check="${orig_check:1}"
        fi
        if [[ "$orig_check" == ~* ]]; then
            # Map the tilde-prefix the user entered back onto the filesystem
            # path produced by compgen so the completion output preserves the symbolic form.
            local tilde_prefix="${orig_check%%/*}"
            local expanded_prefix=""
            if [[ "$tilde_prefix" == "~" ]]; then
                expanded_prefix="$HOME"
            elif [[ "$tilde_prefix" == ~+ ]]; then
                expanded_prefix="$PWD"
            elif [[ "$tilde_prefix" == ~- ]]; then
                expanded_prefix="${OLDPWD-}"
            else
                local tilde_user="${tilde_prefix:1}"
                if [[ -n "$tilde_user" ]]; then
                    # Find the user from passwd.
                    expanded_prefix="$(getent passwd "$tilde_user" 2>/dev/null | cut -d: -f6)"
                fi
            fi
            if [[ -n "$expanded_prefix" ]]; then
                for i in "${!pinyin_matched[@]}"; do
                    if [[ "${pinyin_matched[$i]}" == "$expanded_prefix" ]]; then
                        # Exact home directory.
                        pinyin_matched[$i]="$tilde_prefix"
                    elif [[ "${pinyin_matched[$i]}" == "$expanded_prefix"/* ]]; then
                        local suffix="${pinyin_matched[$i]#"$expanded_prefix/"}"
                        # Join path under the user home.
                        pinyin_matched[$i]="$tilde_prefix/$suffix"
                    fi
                done
            fi
        fi

        local current_results_var="current_results"
        eval "local -a $current_results_var=(\"\${$var_name[@]}\")"
        
        # Merge results and remove duplicates
        local -a all_results
        eval "all_results=(\"\${$current_results_var[@]}\" \"\${pinyin_matched[@]}\")"

        declare -A seen
        local -a unique_results=()
        for item in "${all_results[@]}"; do
            if [[ -z "${seen[$item]}" ]]; then
                seen["$item"]=1
                unique_results+=("$item")
            fi
        done

        eval "$var_name=(\"\${unique_results[@]}\")"
    fi
}
