#!/usr/bin/env bash
# Fetch audio from a URL (YouTube, archive.org, LibriVox, podcast feed,
# anything yt-dlp supports) and convert to the 24 kHz mono WAV format
# our data-prep pipeline expects.
#
# Output is suitable as direct input to `examples/audio_to_manifest`.
#
# Requires: yt-dlp, ffmpeg (install via `brew install yt-dlp ffmpeg`).
#
# Usage:
#   scripts/fetch_audio.sh <url> [output_dir]
#
# Examples:
#   # Single LibriVox chapter
#   scripts/fetch_audio.sh \
#       "https://librivox.org/.../alice-in-wonderland-chapter-1.mp3" \
#       /tmp/voice_corpus
#
#   # YouTube talk
#   scripts/fetch_audio.sh \
#       "https://www.youtube.com/watch?v=XXXXXXXXX" \
#       /tmp/voice_corpus
#
#   # Multiple URLs from a file (one per line)
#   while read url; do
#     scripts/fetch_audio.sh "$url" /tmp/voice_corpus
#   done < urls.txt
#
# Tips for good training audio:
#   - Single dominant speaker per clip (use diarize step if mixed)
#   - At least 10-30 minutes total per persona
#   - Studio-quality > phone > YouTube > heavily-compressed
#   - Emotional variety beats raw duration

set -euo pipefail

if [[ $# -lt 1 ]]; then
    echo "usage: $0 <url> [output_dir]" >&2
    exit 1
fi

URL="$1"
OUT_DIR="${2:-./fetched_audio}"

mkdir -p "$OUT_DIR"

# yt-dlp emits the title as the filename. Strip whitespace + special
# chars in a post-rename so the manifest's `wav` paths stay shell-safe.
TMP_DIR="$(mktemp -d)"
trap "rm -rf '$TMP_DIR'" EXIT

echo "→ downloading audio from: $URL"
yt-dlp \
    --extract-audio \
    --audio-format wav \
    --audio-quality 0 \
    --no-playlist \
    --output "$TMP_DIR/%(title).100s.%(ext)s" \
    "$URL"

# Find what we just downloaded.
SRC_WAV="$(find "$TMP_DIR" -maxdepth 1 -name '*.wav' | head -n1)"
if [[ -z "$SRC_WAV" ]]; then
    echo "error: yt-dlp produced no .wav file" >&2
    exit 2
fi

# Slugify the basename: lowercase, replace non-alphanumeric with `_`,
# collapse runs, trim.
BASE="$(basename "$SRC_WAV" .wav)"
SLUG="$(echo "$BASE" \
    | tr '[:upper:]' '[:lower:]' \
    | sed -E 's/[^a-z0-9]+/_/g; s/^_+|_+$//g')"
DST_WAV="$OUT_DIR/${SLUG}.wav"

echo "→ converting to 24 kHz mono 16-bit PCM: $DST_WAV"
ffmpeg -y -loglevel error \
    -i "$SRC_WAV" \
    -ac 1 -ar 24000 -sample_fmt s16 \
    "$DST_WAV"

# Print duration so the user sees what they got.
DURATION="$(ffmpeg -i "$DST_WAV" 2>&1 | grep Duration | awk '{print $2}' | tr -d ',')"
SIZE_MB="$(du -m "$DST_WAV" | awk '{print $1}')"
echo
echo "✓ wrote $DST_WAV  (${DURATION}, ${SIZE_MB} MB)"
echo
echo "Next step:"
echo "  cargo run -p rtx-csm --release --features metal --example audio_to_manifest -- \\"
echo "    --in $DST_WAV \\"
echo "    --out-dir ${OUT_DIR}_manifest \\"
echo "    --wavlm-sv-weights /tmp/wavlm_sv.safetensors \\"
echo "    --auto-emotion-tag --use-emotion2vec --stage audiobook"