#!/usr/bin/env bash
# Hearthwave — CPU-only setup (build-from-source Docker path).
# One-time host prep + model staging, then:
#   docker compose -f docker-compose.yml up --build
#
# Usage:
#   bash setup-cpu.sh                                  # small.en + large-v3-turbo + voices
#   bash setup-cpu.sh --model base.en                 # smaller/faster Whisper
#   bash setup-cpu.sh --final-model distil-large-v3   # stage a different final model (CPU)
#   bash setup-cpu.sh --final-model none              # skip the final model (single-pass RX)
#   bash setup-cpu.sh --skip-voices

set -euo pipefail

WHISPER_MODEL="small.en"
# Two-tier final-pass model staged by default (~1.6 GB). whisper_model_final="auto"
# prefers large-v3-turbo, so staging it here is what lets "auto" resolve at all.
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,11p' "$0"; exit 0 ;;
    *) echo "Unknown argument: $1  (try --help)" >&2; exit 1 ;;
  esac
done

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_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"
)

[[ -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=""

ct2_repo_for "$WHISPER_MODEL" >/dev/null
[[ -n "$FINAL_MODEL" ]] && ct2_repo_for "$FINAL_MODEL" >/dev/null

echo "==> Creating directories..."
mkdir -p data/journals data/public data/plugins Voices "Models/STT"

# Seed the example plugins (MeshCore + Meshtastic) on first setup so a fresh
# install has working reference plugins to study / enable.
if [[ -d examples/plugins && -z "$(ls -A data/plugins 2>/dev/null)" ]]; then
  cp -r examples/plugins/. data/plugins/
  echo "==> Seeded example plugins into data/plugins (MeshCore, Meshtastic)."
fi

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

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

# CPU-only backend.
COMPUTE_BACKEND=cpu
EOF
echo "==> .env written (COMPUTE_BACKEND=cpu, PULSE_UID=${PULSE_UID})"

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}"
if [[ -n "$FINAL_MODEL" ]]; then
  echo ""
  fetch_repo "$(ct2_repo_for "$FINAL_MODEL")" "Models/STT/${FINAL_MODEL}"
  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

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 "  CPU setup complete. Start the stack:"
echo "        docker compose -f docker-compose.yml up --build"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"