# Bash completion for the `ml` CLI
# Usage:
#   source /path/to/ml_completion.bash
# or
#   echo 'source /path/to/ml_completion.bash' >> ~/.bashrc

_ml_completions()
{
    local cur prev cmds
    COMPREPLY=()
    cur="${COMP_WORDS[COMP_CWORD]}"
    prev="${COMP_WORDS[COMP_CWORD-1]}"

    # Top-level subcommands
    cmds="init sync queue requeue status monitor cancel prune watch dataset experiment narrative outcome info logs annotate validate compare find export"

    # Global options
    global_opts="--help --verbose --quiet --monitor --json"

    # If completing the subcommand itself
    if [[ ${COMP_CWORD} -eq 1 ]]; then
        COMPREPLY=( $(compgen -W "${cmds} ${global_opts}" -- "${cur}") )
        return 0
    fi

    # Handle global options anywhere
    case "${prev}" in
        --help|--verbose|--quiet|--monitor)
            # No further completion after global flags
            return 0
            ;;
    esac

    case "${COMP_WORDS[1]}" in
        init)
            # No specific arguments for init
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        sync)
            # Complete directories for sync
            COMPREPLY=( $(compgen -d -- "${cur}") )
            ;;
        queue)
            queue_opts="--commit --priority --cpu --memory --gpu --gpu-memory --snapshot-id --snapshot-sha256 --args --note --hypothesis --context --intent --expected-outcome --experiment-group --tags --dry-run --validate --explain --force --mlflow --mlflow-uri --tensorboard --wandb-key --wandb-project --wandb-entity -- ${global_opts}"
            case "${prev}" in
                --priority)
                    COMPREPLY=( $(compgen -W "0 1 2 3 4 5 6 7 8 9 10" -- "${cur}") )
                    ;;
                --cpu|--memory|--gpu)
                    COMPREPLY=( $(compgen -W "0 1 2 4 8 16 32" -- "${cur}") )
                    ;;
                --gpu-memory)
                    COMPREPLY=( $(compgen -W "4 8 16 24 32 48" -- "${cur}") )
                    ;;
                --commit|--snapshot-id|--snapshot-sha256|--args|--note|--hypothesis|--context|--intent|--expected-outcome|--experiment-group|--tags|--mlflow-uri|--wandb-key|--wandb-project|--wandb-entity)
                    # Free-form; no special completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${queue_opts}" -- "${cur}") )
                    else
                        # Suggest common job names (static for now)
                        COMPREPLY=( $(compgen -W "train evaluate deploy test baseline" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        requeue)
            requeue_opts="--name --priority --cpu --memory --gpu --gpu-memory --args -- ${global_opts}"
            case "${prev}" in
                --priority)
                    COMPREPLY=( $(compgen -W "0 1 2 3 4 5 6 7 8 9 10" -- "${cur}") )
                    ;;
                --cpu|--memory|--gpu)
                    COMPREPLY=( $(compgen -W "0 1 2 4 8 16 32" -- "${cur}") )
                    ;;
                --gpu-memory)
                    COMPREPLY=( $(compgen -W "4 8 16 24 32 48" -- "${cur}") )
                    ;;
                --name|--args)
                    # Free-form; no special completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${requeue_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        status)
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        monitor)
            COMPREPLY=( $(compgen -W "${global_opts}" -- "${cur}") )
            ;;
        cancel)
            # Job names often free-form; no special completion
            ;;
        prune)
            case "${prev}" in
                --keep)
                    COMPREPLY=( $(compgen -W "1 2 3 4 5" -- "${cur}") )
                    ;;
                --older-than)
                    COMPREPLY=( $(compgen -W "1h 6h 12h 1d 3d 7d" -- "${cur}") )
                    ;;
                *)
                    COMPREPLY=( $(compgen -W "--keep --older-than ${global_opts}" -- "${cur}") )
                    ;;
            esac
            ;;
        watch)
            # Complete directories for watch
            COMPREPLY=( $(compgen -d -- "${cur}") )
            ;;
        dataset)
            COMPREPLY=( $(compgen -W "list register info search verify" -- "${cur}") )
            ;;
        experiment)
            COMPREPLY=( $(compgen -W "log show list" -- "${cur}") )
            ;;
        narrative)
            narrative_opts="set --hypothesis --context --intent --expected-outcome --parent-run --experiment-group --tags --base --json --help"
            case "${prev}" in
                set)
                    COMPREPLY=( $(compgen -W "${narrative_opts}" -- "${cur}") )
                    ;;
                --hypothesis|--context|--intent|--expected-outcome|--parent-run|--experiment-group|--tags|--base)
                    # Free-form completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${narrative_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        outcome)
            outcome_opts="set --outcome --summary --learning --next-step --validation-status --surprise --base --json --help"
            case "${prev}" in
                set)
                    COMPREPLY=( $(compgen -W "${outcome_opts}" -- "${cur}") )
                    ;;
                --outcome)
                    COMPREPLY=( $(compgen -W "validates refutes inconclusive partial" -- "${cur}") )
                    ;;
                --validation-status)
                    COMPREPLY=( $(compgen -W "validates refutes inconclusive" -- "${cur}") )
                    ;;
                --summary|--learning|--next-step|--surprise|--base)
                    # Free-form completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${outcome_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        info|logs|annotate|validate)
            # These commands take a path/id and various options
            info_opts="--base --json --help"
            case "${prev}" in
                --base)
                    COMPREPLY=( $(compgen -d -- "${cur}") )
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${info_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        compare)
            compare_opts="--json --csv --all --fields --help"
            case "${prev}" in
                --fields)
                    COMPREPLY=( $(compgen -W "narrative,metrics,metadata,outcome" -- "${cur}") )
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${compare_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        find)
            find_opts="--json --csv --limit --tag --outcome --dataset --experiment-group --author --after --before --help"
            case "${prev}" in
                --limit)
                    COMPREPLY=( $(compgen -W "10 20 50 100" -- "${cur}") )
                    ;;
                --outcome)
                    COMPREPLY=( $(compgen -W "validates refutes inconclusive partial" -- "${cur}") )
                    ;;
                --tag|--dataset|--experiment-group|--author|--after|--before)
                    # Free-form completion
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${find_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;
        export)
            export_opts="--bundle --anonymize --anonymize-level --base --json --help"
            case "${prev}" in
                --anonymize-level)
                    COMPREPLY=( $(compgen -W "metadata-only full" -- "${cur}") )
                    ;;
                --bundle|--base)
                    COMPREPLY=( $(compgen -f -- "${cur}") )
                    ;;
                *)
                    if [[ "${cur}" == --* ]]; then
                        COMPREPLY=( $(compgen -W "${export_opts}" -- "${cur}") )
                    fi
                    ;;
            esac
            ;;

    return 0
}

complete -F _ml_completions ml