#!/bin/sh
set -eu

usage() {
    cat <<'EOF'
Usage: monarch-kernel [--diagnose] [application arguments]

Environment:
  MONARCH_RUNTIME=auto|cpu|cuda   Select the packaged runtime (default: auto)
  MONARCH_GGUF_DEVICE=auto|cpu|gpu
                                  Override llama.cpp device selection
  MONARCH_NVIDIA_UVM_DEVICE=PATH  UVM device node (default: /dev/nvidia-uvm)
  MONARCH_LAUNCH_QUIET=1          Hide launcher selection messages
EOF
}

case "${1:-}" in
    --help|-h)
        usage
        exit 0
        ;;
esac

launcher_path=$0
while [ -L "$launcher_path" ]; do
    link_target=$(readlink "$launcher_path")
    case "$link_target" in
        /*) launcher_path=$link_target ;;
        *) launcher_path=$(dirname -- "$launcher_path")/$link_target ;;
    esac
done
script_dir=$(CDPATH= cd -- "$(dirname -- "$launcher_path")" && pwd -P)
prefix=$(CDPATH= cd -- "$script_dir/.." && pwd -P)
runtime_dir=${MONARCH_RUNTIME_DIR:-"$prefix/lib/monarch-kernel"}
share_dir=${MONARCH_SHARE_DIR:-"$prefix/share/monarch-kernel"}
cpu_binary="$runtime_dir/monarch-kernel-gguf-cpu"
cuda_binary="$runtime_dir/monarch-kernel-gguf-cuda"
uvm_device=${MONARCH_NVIDIA_UVM_DEVICE:-/dev/nvidia-uvm}
requested=${MONARCH_RUNTIME:-auto}

nvidia_detected=0
if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
    nvidia_detected=1
fi

uvm_ready=0
if [ -c "$uvm_device" ]; then
    uvm_ready=1
fi

cuda_dependencies_ready=0
if [ -x "$cuda_binary" ]; then
    cuda_dependencies_ready=1
    if command -v ldd >/dev/null 2>&1 && ldd "$cuda_binary" 2>&1 | grep -q 'not found'; then
        cuda_dependencies_ready=0
    fi
fi

case "$requested" in
    auto)
        if [ "$cuda_dependencies_ready" -eq 1 ] && [ "$nvidia_detected" -eq 1 ] \
            && [ "$uvm_ready" -eq 1 ]; then
            selected=cuda
            binary=$cuda_binary
        else
            selected=cpu
            binary=$cpu_binary
        fi
        ;;
    cpu)
        selected=cpu
        binary=$cpu_binary
        ;;
    cuda)
        if [ ! -x "$cuda_binary" ]; then
            printf '%s\n' "Monarch: CUDA runtime is not installed in $runtime_dir" >&2
            exit 66
        fi
        if [ "$uvm_ready" -ne 1 ]; then
            cat >&2 <<EOF
Monarch: CUDA was requested, but $uvm_device is unavailable.
On Debian, load it from a privileged host shell with:
  sudo /sbin/modprobe nvidia-current-uvm
  nvidia-modprobe -u -c=0
EOF
            exit 78
        fi
        if [ "$cuda_dependencies_ready" -ne 1 ]; then
            printf '%s\n' "Monarch: CUDA runtime has unresolved shared-library dependencies." >&2
            ldd "$cuda_binary" >&2 || true
            exit 69
        fi
        selected=cuda
        binary=$cuda_binary
        ;;
    *)
        printf '%s\n' "Monarch: invalid MONARCH_RUNTIME=$requested (use auto, cpu, or cuda)" >&2
        exit 64
        ;;
esac

if [ ! -x "$binary" ]; then
    printf '%s\n' "Monarch: selected runtime is missing or not executable: $binary" >&2
    exit 66
fi

if [ -z "${MONARCH_GGUF_DEVICE+x}" ]; then
    if [ "$requested" = cuda ]; then
        MONARCH_GGUF_DEVICE=gpu
    elif [ "$selected" = cuda ]; then
        MONARCH_GGUF_DEVICE=auto
    else
        MONARCH_GGUF_DEVICE=cpu
    fi
    export MONARCH_GGUF_DEVICE
fi

packaged_m2v="$share_dir/models/potion-base-8M"
if [ -z "${MONARCH_M2V_DIR+x}" ] && [ -f "$packaged_m2v/model.safetensors" ]; then
    MONARCH_M2V_DIR=$packaged_m2v
    export MONARCH_M2V_DIR
fi

packaged_ocr="$share_dir/models/ocrs"
if [ -z "${MONARCH_OCR_DIR+x}" ] \
    && [ -f "$packaged_ocr/text-detection-ssfbcj81.rten" ] \
    && [ -f "$packaged_ocr/text-rec-checkpoint-s52qdbqt.rten" ]; then
    MONARCH_OCR_DIR=$packaged_ocr
    export MONARCH_OCR_DIR
fi

if [ "$requested" = auto ] && [ "$selected" = cpu ] && [ -x "$cuda_binary" ] \
    && [ "$nvidia_detected" -eq 1 ] && [ "$uvm_ready" -ne 1 ]; then
    cat >&2 <<EOF
Monarch: NVIDIA was detected, but $uvm_device is unavailable; using CPU.
Run 'sudo /sbin/modprobe nvidia-current-uvm' and 'nvidia-modprobe -u -c=0'
from a privileged host shell to enable CUDA on the next launch.
EOF
fi

if [ "$requested" = auto ] && [ "$selected" = cpu ] && [ -x "$cuda_binary" ] \
    && [ "$nvidia_detected" -eq 1 ] && [ "$uvm_ready" -eq 1 ] \
    && [ "$cuda_dependencies_ready" -ne 1 ]; then
    printf '%s\n' "Monarch: CUDA dependencies are unresolved; using CPU." >&2
fi

if [ "${1:-}" = --diagnose ]; then
    dependency_state=unchecked
    if command -v ldd >/dev/null 2>&1; then
        if ldd "$binary" 2>&1 | grep -q 'not found'; then
            dependency_state=missing
        else
            dependency_state=ok
        fi
    fi
    model_state=unset
    if [ -n "${MONARCH_GGUF_MODEL:-}" ]; then
        if [ -f "$MONARCH_GGUF_MODEL" ]; then
            model_state=present
        else
            model_state=missing
        fi
    fi
    printf '%s\n' \
        "monarch.runtime.requested=$requested" \
        "monarch.runtime.selected=$selected" \
        "monarch.runtime.binary=$binary" \
        "monarch.runtime.dependencies=$dependency_state" \
        "monarch.nvidia.detected=$nvidia_detected" \
        "monarch.nvidia.uvm_device=$uvm_device" \
        "monarch.nvidia.uvm_ready=$uvm_ready" \
        "monarch.cuda.dependencies_ready=$cuda_dependencies_ready" \
        "monarch.gguf.device=$MONARCH_GGUF_DEVICE" \
        "monarch.gguf.model=$model_state" \
        "monarch.m2v.dir=${MONARCH_M2V_DIR:-unset}" \
        "monarch.ocr.dir=${MONARCH_OCR_DIR:-unset}"
    exit 0
fi

if [ "${MONARCH_LAUNCH_QUIET:-0}" != 1 ]; then
    printf '%s\n' "Monarch: launching $selected runtime ($MONARCH_GGUF_DEVICE)" >&2
fi

exec "$binary" "$@"
