#!/usr/bin/env bash
# cmux grok wrapper - installs cmux hooks and captures launch metadata.

find_real_grok() {
    local self_dir
    self_dir="$(cd "$(dirname "$0")" && pwd)"
    local names=(grok grok-macos-aarch64 grok-macos-aarch)
    local IFS=:
    for d in $PATH; do
        [[ "$d" == "$self_dir" ]] && continue
        local name candidate
        for name in "${names[@]}"; do
            candidate="$d/$name"
            [[ -x "$candidate" ]] || continue
            if [[ "$candidate" -ef "$0" ]]; then
                continue
            fi
            printf '%s' "$candidate"
            return 0
        done
    done
    return 1
}

resolve_cmux_bin() {
    local bundled_cli self_dir
    bundled_cli="${CMUX_BUNDLED_CLI_PATH:-}"
    if [[ -n "$bundled_cli" && -x "$bundled_cli" ]]; then
        printf '%s' "$bundled_cli"
        return 0
    fi

    self_dir="$(cd "$(dirname "$0")" && pwd)"
    bundled_cli="$self_dir/cmux"
    if [[ -x "$bundled_cli" ]]; then
        printf '%s' "$bundled_cli"
        return 0
    fi

    bundled_cli="$(command -v cmux || true)"
    if [[ -n "$bundled_cli" ]]; then
        printf '%s' "$bundled_cli"
        return 0
    fi

    return 1
}

cmux_socket_available() {
    local socket="${CMUX_SOCKET_PATH:-}"
    [[ -n "$socket" && -S "$socket" ]] || return 1

    local cmux_bin
    cmux_bin="$(resolve_cmux_bin)" || return 1
    CMUXTERM_CLI_RESPONSE_TIMEOUT_SEC=0.75 \
        "$cmux_bin" --socket "$socket" ping >/dev/null 2>&1
}

encode_launch_argv() {
    {
        printf '%s\0' "$REAL_GROK"
        if (( $# > 0 )); then
            printf '%s\0' "$@"
        fi
    } | base64 | tr -d '\n'
}

grok_option_consumes_value() {
    case "$1" in
        --cwd|--model|-m|--permission-mode|--prompt-file|--prompt-json|\
        --sandbox|--settings|--config|--profile|--resume|-r)
            return 0
            ;;
    esac
    return 1
}

grok_passthrough_option_flag() {
    case "$1" in
        --help|-h|--version|-v|--single|--single=*|-p|--prompt-file|\
        --prompt-file=*|--prompt-json|--prompt-json=*|--output-format|\
        --output-format=*|--best-of-n|--best-of-n=*)
            return 0
            ;;
    esac
    return 1
}

grok_builtin_command_name() {
    case "$1" in
        auth|completion|completions|config|help|login|logout|models|sessions|\
        setup|update|upgrade|version)
            return 0
            ;;
    esac
    return 1
}

grok_command_like_invocation() {
    local candidate="$1"
    [[ -n "$candidate" ]] || return 1
    case "$candidate" in
        *[![:alnum:]_-]*)
            return 1
            ;;
    esac
    grok_builtin_command_name "$candidate"
}

should_install_grok_hooks() {
    (( $# == 0 )) && return 0

    local arg
    local skip_next=false
    for arg in "$@"; do
        if [[ "$skip_next" == true ]]; then
            skip_next=false
            continue
        fi
        case "$arg" in
            --)
                return 0
                ;;
            --resume=*|-r=*)
                return 0
                ;;
            -*)
                if grok_passthrough_option_flag "$arg"; then
                    return 1
                fi
                if [[ "$arg" != *=* ]] && grok_option_consumes_value "$arg"; then
                    skip_next=true
                fi
                continue
                ;;
            *)
                if grok_command_like_invocation "$arg"; then
                    return 1
                fi
                return 0
                ;;
        esac
    done

    return 0
}

REAL_GROK="$(find_real_grok)" || { echo "Error: grok not found in PATH" >&2; exit 127; }

IN_CMUX=0
if [[ -n "${CMUX_SURFACE_ID:-}" ]]; then
    IN_CMUX=1
fi

if [[ "$IN_CMUX" == "1" && "${CMUX_GROK_HOOKS_DISABLED:-}" != "1" ]] \
    && cmux_socket_available \
    && should_install_grok_hooks "$@"; then
    CMUX_BIN="$(resolve_cmux_bin || true)"
    if [[ -n "$CMUX_BIN" ]]; then
        CMUXTERM_CLI_RESPONSE_TIMEOUT_SEC=2 \
            "$CMUX_BIN" --socket "$CMUX_SOCKET_PATH" hooks grok install --yes >/dev/null 2>&1 || true
    fi
    export CMUX_AGENT_LAUNCH_KIND="grok"
    export CMUX_AGENT_LAUNCH_EXECUTABLE="$REAL_GROK"
    export CMUX_AGENT_LAUNCH_ARGV_B64="$(encode_launch_argv "$@")"
    export CMUX_AGENT_LAUNCH_CWD="$PWD"
fi

exec "$REAL_GROK" "$@"
