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

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
models_dir="$repo_root/models"
parakeet_dir="$models_dir/parakeet-tdt-v2"
archive_url="https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/sherpa-onnx-nemo-parakeet-tdt-0.6b-v2-int8.tar.bz2"
vad_url="https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/silero_vad_v5.onnx"

command -v curl >/dev/null || { echo "curl is required" >&2; exit 1; }
command -v tar >/dev/null || { echo "tar is required" >&2; exit 1; }

mkdir -p "$models_dir"
tmp_dir="$(mktemp -d)"
trap 'rm -rf "$tmp_dir"' EXIT

echo "Downloading Parakeet TDT v2 int8 model..."
curl --fail --location --progress-bar "$archive_url" --output "$tmp_dir/parakeet.tar.bz2"
tar -xjf "$tmp_dir/parakeet.tar.bz2" -C "$tmp_dir"

source_dir="$(find "$tmp_dir" -mindepth 1 -maxdepth 1 -type d -name 'sherpa-onnx-nemo-parakeet-tdt-0.6b-v2-int8' -print -quit)"
if [[ -z "$source_dir" ]]; then
  echo "Downloaded archive did not contain the expected Parakeet directory" >&2
  exit 1
fi

rm -rf "$parakeet_dir"
mkdir -p "$parakeet_dir"
cp "$source_dir/encoder.int8.onnx" "$parakeet_dir/"
cp "$source_dir/decoder.int8.onnx" "$parakeet_dir/"
cp "$source_dir/joiner.int8.onnx" "$parakeet_dir/"
cp "$source_dir/tokens.txt" "$parakeet_dir/"

echo "Downloading Silero VAD..."
curl --fail --location --progress-bar "$vad_url" --output "$models_dir/silero_vad_v5.onnx"

echo "Models installed in $models_dir"