#!/usr/bin/env bash
# GPU 1 dedicated to Moss TTS (STT moved to GPU0 with the LLM).
set -euo pipefail
DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$DIR/common.sh"

MODEL="${MODEL:-Rabe3/Moss-Saudi-3}"
PORT="${PORT:-8003}"
# Dedicated GPU1. Colocated Moss needs AR + codec headroom — 0.85 + codec
# reserve OOMs; 0.72 leaves ~28GB used and still beats the old 0.50+STT share.
MEM_FRACTION="${MEM_FRACTION:-0.72}"
MAX_REQ="${MAX_REQ:-24}"
LOG_DIR="${LOG_DIR:-$HOME/agent-logs}"
mkdir -p "$LOG_DIR"

ensure_toolchain
# Same ~/sglang-env as the LLM -- omni is not on PyPI; installed via uv.
ensure_sglang_omni

export CUDA_VISIBLE_DEVICES=1

if ! command -v sgl-omni >/dev/null 2>&1; then
  echo "[tts] FATAL: sgl-omni command not found after install."
  echo "  See https://github.com/sgl-project/sglang-omni/blob/main/docs/get_started/installation.md"
  exit 1
fi

echo "[tts] launching $MODEL on :$PORT [GPU1 dedicated, mem $MEM_FRACTION, max_req $MAX_REQ]"
CONFIG="${CONFIG:-$DIR/moss_tts_local_blackwell.yaml}"
# Avoid interactive transformers trust_remote_code prompt when stdin is not a TTY.
yes y | sgl-omni serve \
  --model-path "$MODEL" \
  --config "$CONFIG" \
  --allowed-media-domain huggingface.co \
  --allowed-media-domain cas-bridge.xethub.hf.co \
  --host 0.0.0.0 \
  --port "$PORT" \
  --mem-fraction-static "$MEM_FRACTION" \
  --max-running-requests "$MAX_REQ" \
  > "$LOG_DIR/tts.log" 2>&1 &
TTS_PID=$!
# Prefer the actual sgl-omni python process over the `yes` pipe leader if needed
sleep 1
OMNI_PID="$(pgrep -f "sgl-omni serve --model-path" | head -1 || true)"
if [ -n "$OMNI_PID" ]; then TTS_PID="$OMNI_PID"; fi

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

echo "$TTS_PID" > /tmp/tts.pid
echo "[tts] up. pid=$TTS_PID"