# Source the official Docker bash completion for basic docker commands
source /usr/share/bash-completion/completions/docker
# Remove all ':' from COMP_WORDBREAKS so Bash treats Docker image names with registry and tag as a single word.
COMP_WORDBREAKS=${COMP_WORDBREAKS//:/}
export COMP_WORDBREAKS

_docker-run() {
    # Save and enable extglob for advanced pattern matching
    local previous_extglob_setting=$(shopt -p extglob)
    shopt -s extglob

    # Initialize completion variables
    COMPREPLY=()
    local cur prev words cword
    _get_comp_words_by_ref -n : cur prev words cword 

    # Find the main command and its position
    local command='run' command_pos=0 subcommand_pos
    local counter=1
    while [ "$counter" -lt "$cword" ]; do
        case "${words[$counter]}" in
            run)
                ;;        # Continue, don't return
            -*)
                ;;        # Skip flags
            =)
                (( counter++ ))  # Skip '='
                ;;
            *)
                command="${words[$counter]}"
                command_pos=$counter
                break  # Found the command
                ;;
        esac
        (( counter++ ))
    done

    # Get list of docker images (with and without tag)
    local docker_images=($(docker images --format "{{.Repository}}:{{.Tag}}"))
    local docker_images_no_tag=($(docker images --format "{{.Repository}}"))

    # Suggest running container names after --name
    if [[ ${prev} == "--name" ]]; then
        COMPREPLY=$(docker ps --format '{{.Names}}')
        COMPREPLY=($(compgen -W "${COMPREPLY[*]}" -- "${cur}"))

    # Suggest image names after --image
    elif [[ ${prev} == "--image" ]]; then
        COMPREPLY=($(compgen -W "${docker_images[*]}" -- "${cur}"))

    # Suggest default command for image with tag
    elif [[ " ${docker_images[*]} " =~ " ${prev} " ]]; then
        # If user types '-' after image, only show custom flags
        if [[ ${cur} == -* ]]; then
            COMPREPLY=($(compgen -W "--help --loc --mwd --mws --no-gpu --no-it --no-name --no-rm --no-tz --no-user --no-x11 --verbose --version" -- "${cur}"))
            return 0
        fi
        # Otherwise show default command
        COMPREPLY=$(docker inspect --format='{{range .Config.Env}}{{if eq (index (split . "=") 0) "DEFAULT_CMD"}}{{index (split . "=") 1}}{{end}}{{end}}' $prev)
        if [[ -z "$COMPREPLY" ]]; then
            COMPREPLY=$(docker inspect --format='{{join .Config.Cmd " "}}' $prev)
        fi

    # Suggest default command for image without tag (latest)
    elif [[ " ${docker_images_no_tag[*]} " =~ " ${prev} " ]] && [[ " ${docker_images[*]} " =~ " ${prev}:latest " ]]; then
        # If user types '-' after image, only show custom flags
        if [[ ${cur} == -* ]]; then
            COMPREPLY=($(compgen -W "--help --loc --mwd --mws --no-gpu --no-it --no-name --no-rm --no-tz --no-user --no-x11 --verbose --version" -- "${cur}"))
            return 0
        fi
        # Otherwise show default command
        COMPREPLY=$(docker inspect --format='{{range .Config.Env}}{{if eq (index (split . "=") 0) "DEFAULT_CMD"}}{{index (split . "=") 1}}{{end}}{{end}}' $prev:latest)
        if [[ -z "$COMPREPLY" ]]; then
            COMPREPLY=$(docker inspect --format='{{join .Config.Cmd " "}}' $prev:latest)
        fi

    # Handle all other cases - normal Docker completion
    else
        # Check if we already have an image in the command
        local has_image=false
        for word in "${words[@]:1}"; do
            if [[ " ${docker_images[*]} " =~ " ${word} " ]] || [[ " ${docker_images_no_tag[*]} " =~ " ${word} " ]]; then
                has_image=true
                break
            fi
        done

        # Only suggest custom flags after image
        if [[ "$has_image" == "true" && ${cur} == -* ]]; then
            COMPREPLY=($(compgen -W "--help --loc --mwd --mws --no-gpu --no-it --no-name --no-rm --no-tz --no-user --no-x11 --verbose --version" -- "${cur}"))
            return 0
        # If current word starts with '-', suggest flags (when no image present)
        elif [[ ${cur} == -* ]]; then
            local docker_run_completions
            # Get official docker run flags
            docker_run_completions=$(docker __complete run -- 2>/dev/null | grep -E "^--" | awk '{print $1}' | grep -v "^:")
            if [[ -n "$docker_run_completions" ]]; then
                COMPREPLY=($(compgen -W "$docker_run_completions" -- "${cur}"))
            fi
            # Add custom flags
            COMPREPLY+=(
                "--help"
                "--image"  
                "--loc"
                "--mwd"
                "--mws"
                "--no-gpu"
                "--no-it"
                "--no-name"
                "--no-rm"
                "--no-tz"
                "--no-user"
                "--no-x11"
                "--verbose"
                "--version"
            )
            # Filter suggestions based on current input
            COMPREPLY=($(compgen -W "${COMPREPLY[*]}" -- "${cur}"))

        # Handle all other cases - normal Docker completion
        else
            # Build filtered_words (official ones) for docker __complete:
            local filtered_words=()
            local skip_next=false
            for word in "${words[@]:1}"; do
                if [[ "$skip_next" == "true" ]]; then
                    skip_next=false
                    continue
                fi
                case "$word" in
                    --image|--name)
                        skip_next=true
                        ;;
                    # Filter out custom flags that Docker doesn't know about
                    --help|--loc|--mwd|--mws|--no-gpu|--no-it|--no-name|--no-rm|--no-tz|--no-user|--no-x11|--verbose|--version)
                        continue  # Skip custom flags
                        ;;
                    *)
                        filtered_words+=("$word")
                        ;;
                esac
            done

            # Get docker run completions for remaining words 
            local docker_run_completions
            docker_run_completions=$(docker __complete run "${filtered_words[@]}" 2>/dev/null | grep -v "^:" | grep -v "Completion ended" | grep -v "\[Debug\]" | grep -v "\[Error\]")
            if [[ -n "$docker_run_completions" ]]; then
                COMPREPLY=($(compgen -W "$docker_run_completions" -- "${cur}"))
            fi
        fi
    fi

    # Restore previous extglob setting
    eval "$previous_extglob_setting"
    return 0
}

# Register the completion function for docker-run
complete -F _docker-run docker-run