#!/usr/bin/env bash
# Hearthwave — AMD ROCm setup (build-from-source Docker path).
# One-time host prep + model staging, then:
#   docker compose -f docker-compose.rocm.yml up --build
#
# Runs the STT final pass on an AMD GPU (ROCm); streaming stays on CPU.
#
# Usage:
#   bash setup-rocm.sh                                  # small.en + large-v3-turbo final + voices
#   bash setup-rocm.sh --model base.en                 # smaller streaming model
#   bash setup-rocm.sh --final-model distil-large-v3   # choose the GPU final-pass model
#   bash setup-rocm.sh --final-model none              # skip the final model (single-pass RX)
#   bash setup-rocm.sh --skip-voices                   # skip Piper voices

set -euo pipefail

WHISPER_MODEL="small.en"
# GPU final pass default. Matches the whisper_model_final="auto" preference
# order (large-v3-turbo > distil-large-v3 > large-v3), so a fresh install has
# the model "auto" actually wants. Staged twice: CT2 (~1.6 GB, CPU fallback)
# and HF transformers (~1.6 GB, GPU). On the Radeon 680M the design-doc
# shootout measured distil-large-v3 ~40% faster on long utterances at a WER
# wash — pass --final-model distil-large-v3 if you want that trade instead.
FINAL_MODEL="large-v3-turbo"
SKIP_VOICES=false

while [[ $# -gt 0 ]]; do
  case "$1" in
    --model)       [[ $# -lt 2 ]] && { echo "Error: --model needs a value." >&2; exit 1; }; WHISPER_MODEL="$2"; shift 2 ;;
    --final-model) [[ $# -lt 2 ]] && { echo "Error: --final-model needs a value." >&2; exit 1; }; FINAL_MODEL="$2"; shift 2 ;;
    --skip-voices) SKIP_VOICES=true; shift ;;
    -h|--help) sed -n '2,13p' "$0"; exit 0 ;;
    *) echo "Unknown argument: $1  (try --help)" >&2; exit 1 ;;
  esac
done

# CT2 (faster-whisper) repo for a streaming/CPU-fallback model.
ct2_repo_for() {
  case "$1" in
    tiny.en) echo "Systran/faster-whisper-tiny.en" ;;
    base.en) echo "Systran/faster-whisper-base.en" ;;
    small.en) echo "Systran/faster-whisper-small.en" ;;
    medium.en) echo "Systran/faster-whisper-medium.en" ;;
    large-v3) echo "Systran/faster-whisper-large-v3" ;;
    distil-large-v3) echo "Systran/faster-distil-whisper-large-v3" ;;
    large-v3-turbo) echo "deepdml/faster-whisper-large-v3-turbo-ct2" ;;
    *) echo "Error: unknown model '$1'." >&2; exit 1 ;;
  esac
}
# HF transformers repo for the GPU final pass (staged at Models/STT/<name>-hf).
hf_repo_for() {
  case "$1" in
    distil-large-v3) echo "distil-whisper/distil-large-v3" ;;
    large-v3) echo "openai/whisper-large-v3" ;;
    large-v3-turbo) echo "openai/whisper-large-v3-turbo" ;;
    *) echo "Error: '$1' has no HF (GPU) mapping. Use distil-large-v3, large-v3, or large-v3-turbo." >&2; exit 1 ;;
  esac
}

HF_BASE="https://huggingface.co/rhasspy/piper-voices/resolve/main"
VOICES=(
  "en_US-ryan-high|en/en_US/ryan/high/en_US-ryan-high"
  "en_US-amy-medium|en/en_US/amy/medium/en_US-amy-medium"
  "en_US-arctic-medium|en/en_US/arctic/medium/en_US-arctic-medium"
  "en_US-hfc_female-medium|en/en_US/hfc_female/medium/en_US-hfc_female-medium"
  "en_US-kathleen-low|en/en_US/kathleen/low/en_US-kathleen-low"
  "en_US-kristin-medium|en/en_US/kristin/medium/en_US-kristin-medium"
  "en_US-lessac-high|en/en_US/lessac/high/en_US-lessac-high"
  "en_US-libritts-high|en/en_US/libritts/high/en_US-libritts-high"
)

# ── Guards ────────────────────────────────────────────────────────────────────
[[ -f backend/requirements.txt ]] || { echo "Error: run from the Hearthwave repo root." >&2; exit 1; }
command -v docker &>/dev/null || { echo "Error: docker not found in PATH." >&2; exit 1; }
docker info &>/dev/null || { echo "Error: Docker daemon not running or no permission." >&2; exit 1; }

# "none"/"off" (and "") mean: single-pass RX, stage no final model.
[[ "$FINAL_MODEL" == "none" || "$FINAL_MODEL" == "off" ]] && FINAL_MODEL=""

# Validate model choices early.
ct2_repo_for "$WHISPER_MODEL" >/dev/null
[[ -n "$FINAL_MODEL" ]] && { ct2_repo_for "$FINAL_MODEL" >/dev/null; hf_repo_for "$FINAL_MODEL" >/dev/null; }

# ── ROCm host checks ──────────────────────────────────────────────────────────
echo "==> Checking ROCm host prerequisites..."
RENDER_GID="$(getent group render | cut -d: -f3 || true)"
if [[ -z "$RENDER_GID" ]]; then
  echo "  WARNING: no 'render' group on this host. The container needs it for /dev/kfd." >&2
  echo "           Create it or install the amdgpu driver, then re-run. Using fallback GID 992." >&2
  RENDER_GID=992
else
  echo "  render group GID: ${RENDER_GID}"
fi
if [[ ! -e /dev/kfd ]]; then
  echo "  WARNING: /dev/kfd not present — the amdgpu/ROCm kernel driver may not be loaded." >&2
  echo "           The container will fail to start without it. Install ROCm/amdgpu first." >&2
fi
if ! id -nG "$(id -un)" | tr ' ' '\n' | grep -qx render; then
  echo "  NOTE: $(id -un) is not in the 'render' group (only affects host-side GPU tools, not the container)." >&2
fi

# ── 1. Directories + seed config ──────────────────────────────────────────────
echo "==> Creating directories..."
mkdir -p data/journals data/public Voices "Models/STT"

if [[ ! -f data/config.json ]]; then
  if [[ -f data/config.json.example ]]; then
    cp data/config.json.example data/config.json
    echo "==> Seeded data/config.json — set your callsign (currently N0CALL)."
  else
    echo "==> Warning: data/config.json.example missing; create data/config.json manually." >&2
  fi
fi

# ── 2. .env ───────────────────────────────────────────────────────────────────
PULSE_UID="$(id -u)"
cat > .env <<EOF
# Hearthwave Docker Compose environment — generated by setup-rocm.sh
PULSE_UID=${PULSE_UID}

# AMD ROCm backend: STT final pass runs on the GPU.
COMPUTE_BACKEND=rocm

# gfx override so unsupported GPUs (e.g. Radeon 680M / gfx1035) present as gfx1030.
HSA_OVERRIDE_GFX_VERSION=10.3.0

# Host 'render' group GID — the container joins it for /dev/kfd access.
RENDER_GID=${RENDER_GID}
EOF
echo "==> .env written (COMPUTE_BACKEND=rocm, RENDER_GID=${RENDER_GID}, PULSE_UID=${PULSE_UID})"

# ── 3. Models (downloaded via a throwaway container — no host python) ─────────
# snapshot_download a HF repo into a target dir under ./Models.
fetch_repo() {
  local repo="$1" target="$2"
  if [[ -d "$target" && -n "$(ls -A "$target" 2>/dev/null)" ]]; then
    echo "==> ${target}: already present, skipping."
    return 0
  fi
  echo "==> Downloading ${repo} -> ${target} ..."
  mkdir -p "$target"
  docker run --rm -v "$(pwd)/Models:/models" python:3.13-slim bash -c \
    "pip install -q --root-user-action=ignore huggingface_hub && python3 -c \"
from huggingface_hub import snapshot_download
snapshot_download('${repo}', local_dir='/models/${target#Models/}')
print('  done.')\""
}

echo ""
fetch_repo "$(ct2_repo_for "$WHISPER_MODEL")" "Models/STT/${WHISPER_MODEL}"   # streaming (CPU)
if [[ -n "$FINAL_MODEL" ]]; then
  echo ""
  echo "==> Final-pass models for ${FINAL_MODEL}:"
  fetch_repo "$(ct2_repo_for "$FINAL_MODEL")" "Models/STT/${FINAL_MODEL}"        # CPU fallback (CT2)
  fetch_repo "$(hf_repo_for "$FINAL_MODEL")"  "Models/STT/${FINAL_MODEL}-hf"     # GPU (HF transformers)
  echo "  New installs default whisper_model_final=\"auto\" and will pick it up on the next"
  echo "  Listen toggle. If data/config.json has whisper_model_final=\"\" (explicit off), set it"
  echo "  to \"auto\" or \"${FINAL_MODEL}\" to enable the two-tier pass."
fi

# ── 4. Voices ─────────────────────────────────────────────────────────────────
if ! $SKIP_VOICES; then
  echo ""
  echo "==> Piper TTS voices..."
  VOICE_SCRIPT=""
  for entry in "${VOICES[@]}"; do
    stem="${entry%%|*}"; hf_path="${entry##*|}"
    for ext in ".onnx" ".onnx.json"; do
      VOICE_SCRIPT+="
if [ -f '/v/${stem}${ext}' ]; then echo '  ${stem}${ext} — present';
else echo '  ${stem}${ext} ...'; wget -q -O '/v/${stem}${ext}' '${HF_BASE}/${hf_path}${ext}' && echo '  ${stem}${ext} done' || { echo '  FAILED ${stem}${ext}' >&2; rm -f '/v/${stem}${ext}'; }; fi"
    done
  done
  docker run --rm -v "$(pwd)/Voices:/v" alpine sh -c "$VOICE_SCRIPT"
fi

echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
echo "  ROCm setup complete. Start the stack:"
echo "        docker compose -f docker-compose.rocm.yml up --build"
echo ""
echo "  NOTE: the first build pulls ~3-4 GB of ROCm torch wheels."
echo "  If /dev/kfd is missing or the GPU is unsupported, the final pass"
echo "  auto-falls back to CPU (set stt_final_device=cpu to force it)."
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"