#!/usr/bin/env bash
# run.sh — one-command launcher for Marcus on the Jetson.
#
# What it does (safely, in order):
#   1. Ensures systemd Ollama is running (with drop-in flags).
#   2. Pulls the VLM model if it's missing.
#   3. Runs the VLM warmup (asks you to put the robot in squat first).
#   4. Health-gates iGPU placement and system memory.
#   5. Pauses for you to stand the robot up.
#   6. Launches Marcus.
#
# Fail-fast: any health check failure aborts BEFORE the robot stands.
#
# Usage:
#   ./run.sh                      normal launch
#   ./run.sh --skip-warmup        skip warmup (only safe if already warm)
#   ./run.sh --skip-stand-prompt  don't pause for stand-up confirmation
#
# First time on a fresh Jetson (one-time, not run here):
#   sudo ./install_ollama_jetson.sh

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$SCRIPT_DIR"

SKIP_WARMUP=0
SKIP_STAND_PROMPT=0
for arg in "$@"; do
    case "$arg" in
        --skip-warmup)       SKIP_WARMUP=1 ;;
        --skip-stand-prompt) SKIP_STAND_PROMPT=1 ;;
        -h|--help)
            grep '^#' "$0" | sed 's/^# \{0,1\}//'
            exit 0
            ;;
        *)
            echo "Unknown argument: $arg" >&2
            echo "Use --help for usage." >&2
            exit 2
            ;;
    esac
done

# ── tiny color helpers ──────────────────────────────────────────────
BOLD='\033[1m'; GREEN='\033[32m'; YELLOW='\033[33m'; RED='\033[31m'; DIM='\033[2m'; OFF='\033[0m'
step() { echo -e "\n${BOLD}▶ $1${OFF}"; }
ok()   { echo -e "  ${GREEN}✓${OFF} $1"; }
warn() { echo -e "  ${YELLOW}⚠${OFF} $1"; }
die()  { echo -e "  ${RED}✗${OFF} $1"; exit 1; }

MODEL=$(python3 -c "import json; print(json.load(open('Config/config_Brain.json'))['ollama_model'])")
HOST=$(python3  -c "import json; print(json.load(open('Config/config_Brain.json'))['ollama_host'])")

# ── 1. Ensure Ollama service is up ─────────────────────────────────
step "1/5  Ollama systemd service"
if ! systemctl is-active --quiet ollama; then
    warn "service is not active — starting it"
    sudo systemctl start ollama
    sleep 2
fi
if ! systemctl is-active --quiet ollama; then
    die "failed to start Ollama; check: systemctl status ollama"
fi
ADJUST=$(systemctl show ollama -p OOMScoreAdjust --value)
if [[ "$ADJUST" != "500" ]]; then
    warn "OOMScoreAdjust is '$ADJUST' (expected 500)"
    warn "drop-in may not be installed — run: sudo ./install_ollama_jetson.sh"
fi
ok "service active (OOMScoreAdjust=$ADJUST)"

# ── 2. Ensure the model is pulled ──────────────────────────────────
step "2/5  VLM model present"
if ! ollama list | awk 'NR>1 {print $1}' | grep -qx "$MODEL"; then
    warn "$MODEL not in store — pulling (~2.2 GB, one-time)"
    ollama pull "$MODEL"
fi
ok "$MODEL available in Ollama store"

# ── 3. Warmup (robot-in-squat safety banner lives inside the script) ──
step "3/5  VLM warmup"
if [[ $SKIP_WARMUP -eq 1 ]]; then
    warn "skipped by --skip-warmup"
else
    ./warmup_vlm.sh "$MODEL"
fi

# ── 4. Health gates ────────────────────────────────────────────────
step "4/5  Health check"

# 4a. PROCESSOR placement
PS_LINE=$(ollama ps | awk -v m="$MODEL" '$1==m {for(i=1;i<=NF;i++) if($i ~ /GPU/) {print $(i-1), $i; exit}}')
PROCESSOR=$(ollama ps | awk -v m="$MODEL" 'NR>1 && $1==m {
    for(i=4;i<=NF;i++) if($i=="GPU" || $i=="CPU/GPU" || $i=="CPU") {print $(i-1) " " $i; break}
}')
if [[ -z "$PROCESSOR" ]]; then
    warn "model not loaded — running without prewarm; first vision call will cold-load"
elif echo "$PROCESSOR" | grep -qE "^100% GPU"; then
    ok "placement: $PROCESSOR (ideal)"
elif echo "$PROCESSOR" | grep -qE "CPU/GPU"; then
    SPLIT=$(echo "$PROCESSOR" | awk '{print $1}')
    warn "placement: $PROCESSOR — partial CPU offload will slow vision queries"
    warn "if unacceptable, lower OLLAMA_GPU_OVERHEAD in install_ollama_jetson.sh"
else
    warn "placement: $PROCESSOR"
fi

# 4b. system memory
read -r MEM_USED MEM_AVAIL SWAP_USED <<< "$(free -m | awk '
    /^Mem:/  { used=$3; avail=$7 }
    /^Swap:/ { sused=$3 }
    END { print used, avail, sused }
')"
echo "  ${DIM}mem used: ${MEM_USED} MiB   available: ${MEM_AVAIL} MiB   swap used: ${SWAP_USED} MiB${OFF}"
if (( SWAP_USED > 500 )); then
    die "swap in use (${SWAP_USED} MiB) — memory is already tight, would thrash under load"
fi
if (( MEM_AVAIL < 1000 )); then
    die "only ${MEM_AVAIL} MiB available — Holosoma/Marcus/camera won't have room"
fi
ok "memory healthy"

# ── 5. Stand the robot, then launch ────────────────────────────────
step "5/5  Launch Marcus"
if [[ $SKIP_STAND_PROMPT -eq 0 ]]; then
    echo -e "  ${YELLOW}Robot can now stand up.${OFF}"
    echo -e "  ${DIM}Press ENTER when the robot is standing, or Ctrl-C to abort.${OFF}"
    read -r
fi
echo

exec python3 run_marcus.py