#!/usr/bin/env bash
# GPU 0, colocated with LLM. Leaves GPU 1 exclusively for Moss TTS.
set -euo pipefail
DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$DIR/common.sh"

MODEL="${MODEL:-CohereLabs/cohere-transcribe-arabic-07-2026}"
PORT="${PORT:-8002}"
# Weights ~3.9GiB. Share GPU0 with granite (~0.68); keep STT small.
MEM_UTIL="${MEM_UTIL:-0.18}"
LOG_DIR="${LOG_DIR:-$HOME/agent-logs}"
mkdir -p "$LOG_DIR"

ensure_hf_home

# Prefer a fully-cached local snapshot so gated-repo 401s can't block restart
# when HF_TOKEN isn't in this shell.
if [[ "$MODEL" != /* ]]; then
  local_snap="$(ls -d "${HF_HOME}/hub/models--${MODEL//\//--}/snapshots/"* 2>/dev/null | head -1 || true)"
  if [ -n "$local_snap" ] && [ -e "$local_snap/model.safetensors" -o -e "$local_snap/model.safetensors.index.json" ]; then
    echo "[stt] using local snapshot: $local_snap"
    MODEL="$local_snap"
  elif [ -z "${HF_TOKEN:-}" ]; then
    echo "[stt] WARNING: HF_TOKEN not set. This is a gated model --"
    echo "  export HF_TOKEN=hf_xxx   (after accepting terms in the browser)"
    echo "  or this will fail to download."
  fi
fi

ensure_toolchain
ensure_venv "stt" "vllm==0.19.0" "vllm[audio]" "librosa"

# cohere_asr is newer than the transformers pin that ships with vllm 0.19.0.
if ! python3 -c "from transformers.models.auto.configuration_auto import CONFIG_MAPPING; assert 'cohere_asr' in CONFIG_MAPPING" 2>/dev/null; then
  echo "[bootstrap] upgrading transformers from git for cohere_asr support"
  pip install "git+https://github.com/huggingface/transformers.git" -q
fi

export CUDA_VISIBLE_DEVICES=0
export VLLM_ATTENTION_BACKEND=FLASHINFER
export VLLM_USE_FLASHINFER_SAMPLER=1
export VLLM_FLASHINFER_FORCE_TENSOR_CORES=1

echo "[stt] launching $MODEL on :$PORT [GPU0 w/ LLM, mem $MEM_UTIL]"
vllm serve "$MODEL" \
  --trust-remote-code \
  --host 0.0.0.0 \
  --port "$PORT" \
  --gpu-memory-utilization "$MEM_UTIL" \
  > "$LOG_DIR/stt.log" 2>&1 &
STT_PID=$!

waited=0; max_wait=600
until curl -sf "http://localhost:${PORT}/health" > /dev/null 2>&1; do
  if [ "$waited" -ge "$max_wait" ]; then
    echo "[stt] TIMEOUT"; tail -n 30 "$LOG_DIR/stt.log"; exit 1
  fi
  if ! kill -0 "$STT_PID" 2>/dev/null; then
    echo "[stt] process died"; tail -n 40 "$LOG_DIR/stt.log"; exit 1
  fi
  sleep 2; waited=$((waited + 2))
done

echo "$STT_PID" > /tmp/stt.pid
echo "[stt] up. pid=$STT_PID"