#!/usr/bin/env bash
# End-to-end: URL + target text + target emotion → emotional voice clone.
#
# Pipeline:
#   1. fetch_audio.sh           (yt-dlp + ffmpeg → 24 kHz WAV)
#   2. examples/audio_to_manifest  (diarize + transcribe → manifest.jsonl)
#   3. pick_context.sh          (rank clips by suitability)
#   4. emotional_speech_n.sh    (N-seed picker over the decoder route)
#
# Recipe defaults (Phase 9 findings):
#   - target=decoder, scale=1.0, layers [2,3]
#   - cfg_schedule=linear:3.0:1.0:25
#   - 5-seed roll, lowest-WER picked
#
# All intermediate state caches by URL hash, so re-running with the same
# --workdir reuses fetch + manifest. Emotional generation is the only
# non-cached step (it depends on text + emotion + seeds, not URL).
#
# Usage:
#   scripts/emotional_clone.sh \
#       --url "https://www.youtube.com/watch?v=..." \
#       --text "What you want said." \
#       --emotion happy \
#       --steering-dir /tmp/ravdess_dec_steering \
#       --out /tmp/cloned_emotional.wav \
#       [--workdir /tmp/voice_clone]            # default: /tmp/voice_clone_<ts>
#       [--seeds 7,42,100,123,256]
#       [--target decoder]                       # or backbone
#       [--wavlm-sv /tmp/wavlm_sv.safetensors]

set -euo pipefail

URL=""
TEXT=""
EMOTION=""
STEERING_DIR=""
OUT=""
WORKDIR=""
SEEDS="7,42,100,123,256"
TARGET="decoder"
WAVLM_SV="/tmp/wavlm_sv.safetensors"
KEEP_ALL=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        --url) URL="$2"; shift 2 ;;
        --text) TEXT="$2"; shift 2 ;;
        --emotion) EMOTION="$2"; shift 2 ;;
        --steering-dir) STEERING_DIR="$2"; shift 2 ;;
        --out) OUT="$2"; shift 2 ;;
        --workdir) WORKDIR="$2"; shift 2 ;;
        --seeds) SEEDS="$2"; shift 2 ;;
        --target) TARGET="$2"; shift 2 ;;
        --wavlm-sv) WAVLM_SV="$2"; shift 2 ;;
        --keep-all) KEEP_ALL=true; shift 1 ;;
        -h|--help) sed -n '2,32p' "$0"; exit 0 ;;
        *) echo "unknown arg: $1" >&2; exit 1 ;;
    esac
done

for v in URL TEXT EMOTION STEERING_DIR OUT; do
    if [[ -z "${!v}" ]]; then
        echo "error: --${v,,} is unset" >&2
        echo "usage: $0 --url U --text T --emotion E --steering-dir D --out O" >&2
        exit 1
    fi
done

if [[ ! -f "$WAVLM_SV" ]]; then
    echo "error: wavlm-sv weights not found at $WAVLM_SV" >&2
    exit 1
fi
if [[ ! -d "$STEERING_DIR" ]]; then
    echo "error: steering dir not found: $STEERING_DIR" >&2
    exit 1
fi

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
WORKSPACE_DIR="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
A2M_BIN="$WORKSPACE_DIR/target/release/examples/audio_to_manifest"
N_SCRIPT="$SCRIPT_DIR/emotional_speech_n.sh"
for f in "$A2M_BIN" "$N_SCRIPT"; do
    if [[ ! -x "$f" ]]; then
        echo "error: missing $f" >&2
        exit 1
    fi
done

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

URL_HASH="$(echo -n "$URL" | shasum | awk '{print $1}' | head -c 12)"

# 1. Fetch audio (cached by URL hash).
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: $(basename "$SRC_WAV")"

# 2. Build manifest (cached).
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. Rank + pick top-scoring context clip.
echo "→ step 3/4: ranking context clips"
RANKED="$WORKDIR/ranked_$URL_HASH.tsv"
"$SCRIPT_DIR/pick_context.sh" "$MANIFEST" 999 2>/dev/null > "$RANKED"
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
CTX_WAV="$(echo "$PICKED" | awk -F'\t' '{print $6}')"
CTX_TEXT="$(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): $(basename "$CTX_WAV")"

# 4. Roll N seeds through emotional_speech_n.sh; copy the winner to --out.
echo "→ step 4/4: rolling $(echo "$SEEDS" | tr ',' '\n' | wc -l | tr -d ' ') seeds for emotion='$EMOTION' on target=$TARGET"
EXTRA=()
[[ "$KEEP_ALL" == true ]] && EXTRA+=("--keep-all")
"$N_SCRIPT" \
    --text "$TEXT" \
    --context-wav "$CTX_WAV" \
    --context-text "$CTX_TEXT" \
    --emotion "$EMOTION" \
    --steering-dir "$STEERING_DIR" \
    --target "$TARGET" \
    --seeds "$SEEDS" \
    --wavlm-sv "$WAVLM_SV" \
    --out "$OUT" \
    ${EXTRA[@]+"${EXTRA[@]}"}

echo
echo "✓ wrote $OUT"