#!/usr/bin/env bash

set -euo pipefail

# get_rnnoise_model.sh — fetch an RNNoise model into a local models dir
#
# Prefers known-good rnnoise-nu models. You can override with:
#   RN_URL, RN_TARGET_DIR, RN_TARGET_NAME
#
# Usage:
#   Bash/get_rnnoise_model.sh            # interactive download
#   RN_TARGET_DIR=./models Bash/get_rnnoise_model.sh --yes

# Source common library for shared functions
SCRIPT_DIR="$(dirname "$(readlink -f "$0")")"
# shellcheck source=../../lib/common.sh
source "$SCRIPT_DIR/../../lib/common.sh"

YES=false
while [[ $# -gt 0 ]]; do
  case "$1" in
    -y | --yes)
      YES=true
      shift
      ;;
    *)
      echo "Unknown option: $1" >&2
      exit 2
      ;;
  esac
done

RN_TARGET_DIR=${RN_TARGET_DIR:-"$(dirname "$0")/models"}
RN_TARGET_NAME=${RN_TARGET_NAME:-"rnnoise_model.rnnn"}

mkdir -p "$RN_TARGET_DIR"
dest="$RN_TARGET_DIR/$RN_TARGET_NAME"

if [[ -f $dest ]]; then
  echo "Model already exists at: $dest"
  exit 0
fi

if ! $YES; then
  if ! ask_yes_no "Download RNNoise model to $dest?"; then
    echo "Aborted."
    exit 1
  fi
fi

if ! has_cmd curl && ! has_cmd wget; then
  echo "Error: Need curl or wget to download RNNoise model." >&2
  exit 3
fi

# Helper: try to download a URL to destination, exit 0 on success
# Usage: try_download_model URL DEST
try_download_model() {
  local url="$1"
  local dest="$2"
  local tmp
  tmp=$(mktemp)
  echo "Attempting to download RNNoise model from: $url" >&2
  if has_cmd curl; then
    curl -fsSL "$url" -o "$tmp" 2> /dev/null || true
  else
    wget -qO "$tmp" "$url" 2> /dev/null || true
  fi
  if [[ -s $tmp ]]; then
    mv "$tmp" "$dest"
    echo "Saved RNNoise model to: $dest" >&2
    exit 0
  fi
  rm -f "$tmp" || true
}

# Priority 1: explicit URL
if [[ -n ${RN_URL:-} ]]; then
  echo "Downloading RNNoise model from RN_URL: $RN_URL" >&2
  try_download_model "$RN_URL" "$dest"
  echo "Warning: RN_URL download failed; continuing to fallback sources." >&2
fi

# Priority 2: rnnoise-nu known models (GregorR)
NU_URLS=(
  "https://raw.githubusercontent.com/GregorR/rnnoise-nu/master/src/models/sh.rnnn"
  "https://raw.githubusercontent.com/GregorR/rnnoise-nu/master/src/models/lq.rnnn"
  "https://raw.githubusercontent.com/GregorR/rnnoise-nu/master/src/models/mp.rnnn"
  "https://raw.githubusercontent.com/GregorR/rnnoise-nu/master/src/models/bd.rnnn"
  "https://raw.githubusercontent.com/GregorR/rnnoise-nu/master/src/models/cb.rnnn"
)
for u in "${NU_URLS[@]}"; do
  try_download_model "$u" "$dest"
done

# Priority 2b: arnndn-models fallback (richardpl)
RNNDN_URLS=(
  "https://raw.githubusercontent.com/richardpl/arnndn-models/master/sh.rnnn"
)
for u in "${RNNDN_URLS[@]}"; do
  try_download_model "$u" "$dest"
done

# Priority 3: repo archives (rnnoise-nu and arnndn-models)
ARCHIVES=(
  "https://github.com/GregorR/rnnoise-nu/archive/refs/heads/master.zip"
  "https://github.com/richardpl/arnndn-models/archive/refs/heads/master.zip"
)
for aurl in "${ARCHIVES[@]}"; do
  echo "Attempting to download archive: $aurl" >&2
  tmpdir=$(mktemp -d)
  archive="$tmpdir/models.zip"
  set +e
  if has_cmd curl; then
    curl -fL "$aurl" -o "$archive"
  else
    wget -O "$archive" "$aurl"
  fi
  status=$?
  set -e
  if [[ $status -ne 0 ]]; then
    rm -rf "$tmpdir" || true
    continue
  fi
  if has_cmd bsdtar; then
    bsdtar -xf "$archive" -C "$tmpdir"
  elif has_cmd unzip; then
    unzip -q "$archive" -d "$tmpdir"
  else
    echo "Warning: Need bsdtar or unzip to extract archive; skipping archive method." >&2
    rm -rf "$tmpdir" || true
    continue
  fi
  mapfile -t nnfiles < <(bash -lc 'shopt -s globstar nullglob; for f in '"$tmpdir"'/**/*.rnnn '"$tmpdir"'/**/*.nn; do [[ -f "$f" ]] && echo "$f"; done')
  if [[ ${#nnfiles[@]} -gt 0 ]]; then
    cp -f "${nnfiles[0]}" "$dest"
    echo "Saved RNNoise model to: $dest (from archive)" >&2
    rm -rf "$tmpdir" || true
    exit 0
  fi
  rm -rf "$tmpdir" || true
done

# Priority 4: Arch-based AUR packages and search only .nn/.rnnn
if has_cmd yay; then
  echo "Attempting to install AUR packages that may include RNNoise models..." >&2
  set +e
  yay -S --noconfirm denoiseit-git 2> /dev/null
  yay -S --noconfirm speech-denoiser-git 2> /dev/null
  set -e
  mapfile -t found < <(bash -lc 'shopt -s globstar nullglob; for f in /usr/share/**/*.nn /usr/share/**/*.rnnn /usr/local/share/**/*.nn /usr/local/share/**/*.rnnn; do [[ -f "$f" ]] && echo "$f"; done' 2> /dev/null || true)
  if [[ ${#found[@]} -gt 0 ]]; then
    echo "Found candidate models:" >&2
    printf '  %s\n' "${found[@]}" >&2
    cp -f "${found[0]}" "$dest"
    echo "Copied model to: $dest" >&2
    exit 0
  fi
fi

echo "Error: Could not obtain an RNNoise model automatically." >&2
echo "Hint: Set RN_URL to a reachable model URL, or place a model file at: $dest" >&2
exit 5
