#!/usr/bin/env bash
# End-to-end voice clone: source URL + target text → cloned WAV.
#
# Chains:
#   1. scripts/fetch_audio.sh                 (yt-dlp + ffmpeg → 24kHz WAV)
#   2. examples/audio_to_manifest             (Silero VAD + WavLM-SV diarize
#                                              + Moonshine transcribe → JSONL)
#   3. scripts/pick_context.sh                (rank context clips)
#   4. examples/generate                      (CSM-1B context conditioning)
#
# The picker is a coarse filter (2026-04-29 A/B); this script defaults to
# its top-scoring clip per (source, speaker0). Override with --context-wav
# / --context-text directly to `generate` if you want a different clip.
#
# Usage:
#   scripts/clone_voice.sh \
#       --url "https://www.youtube.com/watch?v=..." \
#       --text "What you want the cloned voice to say." \
#       --out /tmp/cloned.wav \
#       [--workdir /tmp/voice_clone]            # defaults: /tmp/voice_clone_<ts>
#       [--wavlm-sv /tmp/wavlm_sv.safetensors]  # default: /tmp/wavlm_sv.safetensors
#       [--speaker 0]                            # output speaker id (default 0)
#
# The wrapper preserves `workdir` so re-runs can skip refetch / remanifest:
# pass the same --workdir and the existing fetched.wav / manifest.jsonl
# are reused.

set -euo pipefail

URL=""
TEXT=""
OUT=""
WORKDIR=""
WAVLM_SV="/tmp/wavlm_sv.safetensors"
SPEAKER=0

while [[ $# -gt 0 ]]; do
    case "$1" in
        --url) URL="$2"; shift 2 ;;
        --text) TEXT="$2"; shift 2 ;;
        --out) OUT="$2"; shift 2 ;;
        --workdir) WORKDIR="$2"; shift 2 ;;
        --wavlm-sv) WAVLM_SV="$2"; shift 2 ;;
        --speaker) SPEAKER="$2"; shift 2 ;;
        -h|--help)
            sed -n '2,28p' "$0"
            exit 0
            ;;
        *)
            echo "unknown arg: $1" >&2
            exit 1
            ;;
    esac
done

if [[ -z "$URL" || -z "$TEXT" || -z "$OUT" ]]; then
    echo "usage: $0 --url <url> --text <text> --out <wav> [--workdir <dir>] [--wavlm-sv <st>] [--speaker N]" >&2
    exit 1
fi

if [[ ! -f "$WAVLM_SV" ]]; then
    echo "wavlm-sv weights not found at $WAVLM_SV — run examples/wavlm_sv_convert first." >&2
    exit 1
fi

# Locate sibling scripts + cargo binaries by walking up to crate root.
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
CRATE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
WORKSPACE_DIR="$(cd "$CRATE_DIR/../../.." && pwd)"
GEN_BIN="$WORKSPACE_DIR/target/release/examples/generate"
A2M_BIN="$WORKSPACE_DIR/target/release/examples/audio_to_manifest"

for bin in "$GEN_BIN" "$A2M_BIN"; do
    if [[ ! -x "$bin" ]]; then
        echo "binary missing: $bin" >&2
        echo "build with: cargo build -p rtx-csm --release --example generate --example audio_to_manifest" >&2
        exit 1
    fi
done

if [[ -z "$WORKDIR" ]]; then
    WORKDIR="/tmp/voice_clone_$(date +%s)"
fi
mkdir -p "$WORKDIR"
echo "→ workdir: $WORKDIR"

# 1. Fetch audio if we haven't already (cache by URL hash).
URL_HASH="$(echo -n "$URL" | shasum | awk '{print $1}' | head -c 12)"
FETCH_DIR="$WORKDIR/fetch_$URL_HASH"
if [[ ! -d "$FETCH_DIR" ]] || ! ls "$FETCH_DIR"/*.wav >/dev/null 2>&1; then
    mkdir -p "$FETCH_DIR"
    echo "→ step 1/4: fetching audio"
    "$SCRIPT_DIR/fetch_audio.sh" "$URL" "$FETCH_DIR"
else
    echo "→ step 1/4: reusing cached fetch in $FETCH_DIR"
fi
SRC_WAV="$(ls "$FETCH_DIR"/*.wav | head -n1)"
echo "  source: $SRC_WAV"

# 2. Build manifest (skip if it already exists for this source).
MANIFEST_DIR="$WORKDIR/manifest_$URL_HASH"
MANIFEST="$MANIFEST_DIR/manifest.jsonl"
if [[ ! -f "$MANIFEST" ]]; then
    echo "→ step 2/4: diarize + transcribe → manifest"
    "$A2M_BIN" \
        --in "$SRC_WAV" \
        --out-dir "$MANIFEST_DIR" \
        --wavlm-sv-weights "$WAVLM_SV"
else
    echo "→ step 2/4: reusing manifest at $MANIFEST"
fi

# 3. Pick the top-ranked context clip (highest-score row across all groups).
# Materialize to a tempfile so set -o pipefail isn't tripped by `head`'s
# SIGPIPE to upstream `sort`.
echo "→ step 3/4: ranking context clips"
RANKED="$WORKDIR/ranked_$URL_HASH.tsv"
"$SCRIPT_DIR/pick_context.sh" "$MANIFEST" 999 2>/dev/null > "$RANKED"
# `awk 'NR==1'` instead of `head -n 1` so we don't SIGPIPE the upstream
# sort (set -o pipefail would otherwise kill the script with exit 141).
PICKED="$(tail -n +2 "$RANKED" | sort -t $'\t' -k2,2gr | awk 'NR==1')"
if [[ -z "$PICKED" ]]; then
    echo "no scorable clips in manifest — diarization may have failed." >&2
    exit 2
fi
PICKED_WAV="$(echo "$PICKED" | awk -F'\t' '{print $6}')"
PICKED_TXT="$(echo "$PICKED" | awk -F'\t' '{print $7}' | sed 's/^← *//; s/^[[:space:]]*//')"
PICKED_SCORE="$(echo "$PICKED" | awk -F'\t' '{print $2}')"
PICKED_DUR="$(echo "$PICKED" | awk -F'\t' '{print $3}')"
echo "  picked (score=$PICKED_SCORE, dur=${PICKED_DUR}s):"
echo "    $PICKED_WAV"
echo "    \"$(echo "$PICKED_TXT" | head -c 80)...\""

# 4. Generate with the picked context. context_speaker=1 keeps the cloning
# reference distinct from the output speaker slot.
echo "→ step 4/4: generating cloned utterance"
"$GEN_BIN" \
    --text "$TEXT" \
    --speaker "$SPEAKER" \
    --context-wav "$PICKED_WAV" \
    --context-text "$PICKED_TXT" \
    --context-speaker 1 \
    --out "$OUT"

echo
echo "✓ wrote $OUT"