#!/usr/bin/env bash
# Sourced by start_llm.sh / start_stt.sh / start_tts.sh. Not run directly.

SUDO=""
if [ "$(id -u)" -ne 0 ]; then SUDO="sudo"; fi

# Models (~30G for granite+cohere+moss) will not fit next to two CUDA envs on a
# 32G overlay. Keep the HF cache on /dev/shm (RAM-backed; this box has ~2Ti).
ensure_hf_home () {
  local dest="/dev/shm/hf_home"
  local ws="${WORKSPACE:-/workspace}"
  mkdir -p "$dest"
  if [ -d "${ws}/.hf_home" ] && [ ! -L "${ws}/.hf_home" ]; then
    # Migrate any leftover overlay cache once.
    cp -a "${ws}/.hf_home/." "$dest/" 2>/dev/null || true
    rm -rf "${ws}/.hf_home"
  fi
  ln -sfn "$dest" "${ws}/.hf_home"
  export HF_HOME="$dest"
  export HUGGINGFACE_HUB_CACHE="$dest/hub"
  # Persist for login shells / supervisor
  if [ -d "$ws" ] && ! grep -q '^HF_HOME=' "${ws}/.env" 2>/dev/null; then
    printf 'HF_HOME="%s"\nHUGGINGFACE_HUB_CACHE="%s/hub"\n' "$dest" "$dest" >> "${ws}/.env"
  fi
}

# Idempotent: only installs anything if nvcc or its matching g++ is actually missing.
# Vast.ai templates vary a lot -- some ship a complete matched toolchain, some don't.
ensure_toolchain () {
  if ! command -v nvcc >/dev/null 2>&1; then
    echo "[bootstrap] nvcc not found -- installing CUDA toolkit (this can take a few minutes)"
    $SUDO apt-get update -y
    wget -q https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.1-1_all.deb -O /tmp/cuda-keyring.deb
    $SUDO dpkg -i /tmp/cuda-keyring.deb
    $SUDO apt-get update -y
    $SUDO apt-get install -y cuda-toolkit-13-0
  fi

  # nvcc shells out to plain `gcc`, whose version can silently differ from
  # whatever g++ you have -- each gcc major version ships its own cc1plus.
  local gcc_major
  gcc_major="$(gcc -dumpversion | cut -d. -f1)"
  local cc1plus_path
  cc1plus_path="$(gcc -print-prog-name=cc1plus)"
  if [ ! -f "$cc1plus_path" ]; then
    echo "[bootstrap] cc1plus missing for gcc-$gcc_major -- installing g++-$gcc_major"
    $SUDO apt-get update -y
    $SUDO apt-get install -y "g++-${gcc_major}" build-essential
  fi
}

# ensure_venv <name> <pip-install-args...>
# Creates ~/<name>-env if missing, then pip installs the given packages inside it.
# Safe to call every run -- pip install is a no-op if already satisfied.
ensure_venv () {
  local name="$1"; shift
  local venv_dir="$HOME/${name}-env"
  if [ ! -f "$venv_dir/bin/activate" ]; then
    echo "[bootstrap] creating venv: $venv_dir"
    python3 -m venv "$venv_dir"
  fi
  # shellcheck disable=SC1090
  source "$venv_dir/bin/activate"
  # 32G container disks fill fast with 3 CUDA stacks + pip cache.
  export PIP_NO_CACHE_DIR=1
  export UV_NO_CACHE=1
  pip install --upgrade pip -q
  if [ "$#" -gt 0 ]; then
    echo "[bootstrap] pip install: $*"
    pip install "$@" -q
  fi
}

# Shared env for LLM + TTS. sglang-omni is not on PyPI and pins sglang==0.5.12.post1;
# plain pip hits a protobuf conflict (descript-audiotools vs s3prl) that upstream
# solves with uv override-dependencies.
ensure_sglang_omni () {
  local venv_dir="$HOME/sglang-env"
  local src="/dev/shm/sglang-omni-src"
  local overrides="/tmp/sgl-omni-overrides.txt"

  ensure_hf_home
  if [ ! -f "$venv_dir/bin/activate" ]; then
    echo "[bootstrap] creating venv: $venv_dir"
    python3 -m venv "$venv_dir"
  fi
  # shellcheck disable=SC1090
  source "$venv_dir/bin/activate"
  export PIP_NO_CACHE_DIR=1
  export UV_NO_CACHE=1

  if [ ! -d "$src/.git" ]; then
    echo "[bootstrap] cloning sglang-omni -> $src"
    rm -rf "$src"
    git clone --depth 1 https://github.com/sgl-project/sglang-omni.git "$src"
  fi

  # Blackwell (sm_120): force SDPA in Moss vocoder + audio tokenizer.
  _patch_moss_blackwell "$src"

  if command -v sgl-omni >/dev/null 2>&1 && python3 -c "import sglang" 2>/dev/null; then
    return 0
  fi

  # Matches [tool.uv] override-dependencies in upstream pyproject.toml
  printf 'protobuf>=6.31.1,<7.0.0\n' > "$overrides"

  echo "[bootstrap] uv pip install -e sglang-omni (protobuf override)"
  uv pip install --python "$venv_dir/bin/python" -e "$src" --override "$overrides"
}

_patch_moss_blackwell () {
  local src="$1"
  local vocoder="$src/sglang_omni/models/moss_tts_local/vocoder_decoder.py"
  local tokenizer="$src/sglang_omni/models/moss_tts_local/audio_tokenizer.py"

  if [ -f "$vocoder" ] && ! grep -q 'major >= 10' "$vocoder" 2>/dev/null; then
    python3 - "$vocoder" <<'PY'
from pathlib import Path
import sys
path = Path(sys.argv[1])
text = path.read_text()
old = """try:
    from sglang.jit_kernel.flash_attention import flash_attn_varlen_func
except ImportError:
    flash_attn_varlen_func = None
"""
new = """try:
    from sglang.jit_kernel.flash_attention import flash_attn_varlen_func

    # FA3 (sgl-kernel) rejects Blackwell (sm_120). FA4 cute varlen is not API-
    # compatible with the FA2 call site used when FA3 is disabled. Fall back to
    # the dense SDPA path in resolve_attention_implementation().
    if torch.cuda.is_available():
        _major, _ = torch.cuda.get_device_capability()
        if _major >= 10:
            flash_attn_varlen_func = None
except ImportError:
    flash_attn_varlen_func = None
"""
if old not in text:
    raise SystemExit(f"patch target not found in {path}")
path.write_text(text.replace(old, new, 1))
print(f"[bootstrap] patched Moss vocoder for Blackwell: {path}")
PY
  fi

  if [ -f "$tokenizer" ] && ! grep -q 'Blackwell: flash-attn' "$tokenizer" 2>/dev/null; then
    python3 - "$tokenizer" <<'PY'
from pathlib import Path
import sys
path = Path(sys.argv[1])
text = path.read_text()
needle = "    model.eval()\n    model.to(device)\n    return MossTTSLocalAudioTokenizer("
insert = """    model.eval()
    model.to(device)
    # Blackwell: flash-attn FA2/FA3/FA4 paths are broken for this codec; force SDPA.
    if hasattr(model, "set_attention_implementation"):
        model.set_attention_implementation("sdpa")
    for module in model.modules():
        if hasattr(module, "attn_implementation"):
            module.attn_implementation = "sdpa"
        if hasattr(module, "attention_implementation"):
            module.attention_implementation = "sdpa"
    return MossTTSLocalAudioTokenizer("""
if needle not in text:
    raise SystemExit(f"patch target not found in {path}")
path.write_text(text.replace(needle, insert, 1))
print(f"[bootstrap] patched Moss audio tokenizer for Blackwell: {path}")
PY
  fi
}