#!/usr/bin/env bash
set -euo pipefail

MODEL_ID="x-asr-zh-en-960ms"
CHUNK_DIR="chunk-960ms-model"
BASE_URL="https://huggingface.co/GilgameshWind/X-ASR-zh-en/resolve/main/deployment/models/${CHUNK_DIR}"

APP_SUPPORT_DIR="${PRIVATEVOICE_APP_SUPPORT_DIR:-$HOME/Library/Application Support/PrivateVoice Input}"
MODELS_ROOT="${PRIVATEVOICE_MODELS_ROOT:-$APP_SUPPORT_DIR/models}"
FINAL_DIR="$MODELS_ROOT/$MODEL_ID"
RUN_ID="$(date +%Y%m%d.%H%M%S)"
STAGING_DIR="$MODELS_ROOT/.staging/${MODEL_ID}-${RUN_ID}"
BACKUP_DIR="$FINAL_DIR.backup-${RUN_ID}"

FILES=(
  "encoder-960ms.onnx"
  "decoder-960ms.onnx"
  "joiner-960ms.onnx"
  "tokens.txt"
)

cleanup() {
  rm -rf "$STAGING_DIR"
}
trap cleanup EXIT

mkdir -p "$STAGING_DIR"

for file in "${FILES[@]}"; do
  url="$BASE_URL/$file"
  dest="$STAGING_DIR/$file"
  echo "Downloading $file"
  curl -L -sS --fail --retry 5 --retry-delay 2 --continue-at - --output "$dest" "$url"
  if [[ ! -s "$dest" ]]; then
    echo "Downloaded file is empty: $dest" >&2
    exit 1
  fi
  echo "Downloaded $file"
done

xattr -cr "$STAGING_DIR" 2>/dev/null || true

mkdir -p "$MODELS_ROOT"
if [[ -e "$FINAL_DIR" ]]; then
  mv "$FINAL_DIR" "$BACKUP_DIR"
fi

mv "$STAGING_DIR" "$FINAL_DIR"
trap - EXIT

if [[ -e "$BACKUP_DIR" ]]; then
  rm -rf "$BACKUP_DIR"
fi

echo "Installed $MODEL_ID to:"
echo "$FINAL_DIR"
du -sh "$FINAL_DIR" 2>/dev/null || true