# Functions for docker in Lima for macOS.
# Requires lima, docker, docker-compose and docker-credentials-helper to be installed.

if [ "${OS/Darwin}" != "${OS}" ]; then
    LIMA_DOCKER=1

    # Candidates cobra offers for the last argument, one per line, without the directive line or descriptions
    function _docker_complete() {
        local -a out=("${(@f)$(command docker __complete "$@" 2>/dev/null)}")
        print -rl -- "${(@)${(@)out:#:*}%%$'\t'*}"
    }

    function _docker_with_platform() {
        local platform=$1 extra=$2
        shift 2
        local -a args=("$@")

        # Completion requests and explicitly pinned platforms pass through untouched
        if [[ $1 == __complete* ]] || (( ${args[(I)(--platform|--platform=*)]} )); then
            command docker "$@"
            return
        fi

        # Walk the command path; stop at the first command that accepts --platform
        local -i i pos=0
        local -a cmds flags
        for (( i = 1; i <= $#args; i++ )); do
            [[ ${args[i]} == -* ]] && continue

            cmds=("${(@f)$(_docker_complete "${(@)args[1,i-1]}" "")}")
            if (( ! ${cmds[(Ie)${args[i]}]} )); then
                # Not a command: either the previous option's value, or we've left the command path
                (( i > 1 )) && [[ ${args[i-1]} == -* && ${args[i-1]} != *=* ]] && continue
                break
            fi

            flags=("${(@f)$(_docker_complete "${(@)args[1,i]}" -)}")
            if (( ${flags[(Ie)--platform]} )); then
                pos=$i
                break
            fi
        done

        if (( pos )); then
            local -a inject=(--platform=$platform)
            local f name
            for f in ${=extra}; do
                name=${f%%=*}
                # Only if this command accepts it and the user didn't pass it already
                (( ${flags[(Ie)$name]} && ! ${args[(I)$name(|=*)]} )) && inject+=("$f")
            done
            args=("${(@)args[1,pos]}" $inject "${(@)args[pos+1,-1]}")
            print -u2 $'\e[34m-> docker '"${(j: :)${(@q-)args}}"$'\e[0m'
        fi

        command docker "${args[@]}"
    }

    function docker() {
        _docker_with_platform "linux/arm64/v8" "" "$@"
    }

    function docker-amd64() {
        _docker_with_platform "linux/amd64" "--device=lima-vm.io/rosetta=cached" "$@"
    }
fi