# ollama(1) completion                                     -*- shell-script -*-
# vi: ft=sh sw=4
# SPDX-License-Identifier: Apache-2.0
# shellcheck disable=SC2034

_comp_ollama()
{
    local cur prev words cword comp_args
    _comp_initialize -n =: -- "$@" || return
    words=("${words[@]:0:$cword+1}")

    if ((cword == 2)); then
	# Models listing commands.
	case "$prev" in
	    pull)
		local models_list=/usr/share/ollama/models-list.txt
		local tags_list=/usr/share/ollama/tags-list.txt
		if [[ $cur == *:* ]]; then
		    _comp_compgen_split -- "$(LANG=C look "$cur" "$tags_list")"
		    _comp_ltrim_colon_completions "$cur"
		else
		    _comp_compgen_split -- "$(LANG=C look "$cur" "$models_list")"
		    if [[ ${COMPREPLY-} != "$cur" ]]; then
			# Allow user to enter ':' after the model name or a space if
			# another tab is pressed.
			compopt -o nospace
		    fi
		fi
		;;
	    cp|push|rm|run|show)
		_comp_compgen_split -- "$(ollama list 2>/dev/null | tail -n+2 | grep -Eo '^\S+')"
		_comp_ltrim_colon_completions "$cur"
		;;
	    stop)
		_comp_compgen_split -- "$(ollama ps 2>/dev/null | tail -n+2 | grep -Eo '^\S+')"
		_comp_ltrim_colon_completions "$cur"
		;;
	esac
	((${#COMPREPLY[@]})) &&	return
    fi

    # Fallback to completing Cobra built-ins.
    local requestComp lastParam lastChar args
    args=("${words[@]:1}")
    requestComp="${words[0]} __completeNoDesc ${args[*]}"
    stParam=${words[$((${#words[@]}-1))]}
    lastChar=${lastParam:$((${#lastParam}-1)):1}
    if [[ -z $cur && $lastChar != = ]]; then
	requestComp="$requestComp ''"
    fi
    if [[ $cur == -*=* ]]; then
	cur="${cur#*=}"
    fi
    local out directive
    out=$(eval "$requestComp" 2>/dev/null)
    directive=${out##*:}
    out=${out%:*}
    if [[ $directive == "$out" ]]; then
	directive=0
    fi
    local shellCompDirectiveError=1
    local shellCompDirectiveNoSpace=2
    local shellCompDirectiveNoFileComp=4
    local shellCompDirectiveFilterFileExt=8
    local shellCompDirectiveFilterDirs=16
    local shellCompDirectiveKeepOrder=32
    if (((directive & shellCompDirectiveError) != 0)); then
	return
    else
	if (((directive & shellCompDirectiveNoSpace) != 0)); then
	    compopt -o nospace
	fi
	if (((directive & shellCompDirectiveKeepOrder) != 0)); then
	    compopt -o nosort
	fi
	if (((directive & shellCompDirectiveNoFileComp) != 0)); then
	    compopt +o default
	fi
    fi
    _comp_compgen_split -- "$out"
} &&
    complete -F _comp_ollama ollama
