#!/usr/bin/env bash
# Rank manifest clips by suitability as voice-cloning context for
# CSM-1B's --context-wav / --context-text mechanism.
#
# Findings from 2026-04-29 voice-cloning experiments:
#   - context clips < 10 s often produce clipping (output saturates int16)
#   - context clips > ~14 s can blow the model's effective context window
#     and degrade fidelity on the new prompt
#   - clips with healthy RMS (~-15 to -25 dB) and no extreme peaks
#     (peak < ~-3 dBFS) yield clean clones
#   - speaker labels are PER-FILE in our manifests (diarizer restarts);
#     this script preserves that and groups by (source_file, speaker_id)
#
# Usage:
#   scripts/pick_context.sh <manifest.jsonl> [top_n=5]
#
# Outputs (per top group, tab-separated):
#   score   duration_s   rms_db   peak_db   wav_path   transcript
#
# Pipe through `head` to grab the top scorer:
#   scripts/pick_context.sh manifest.jsonl | head -1 | cut -f5
#
# Then feed the wav + transcript to examples/generate:
#   target/release/examples/generate \
#       --text "..." \
#       --context-wav "<picked.wav>" \
#       --context-text "<picked transcript>" \
#       --out output.wav

set -euo pipefail

MANIFEST="${1:?usage: $0 <manifest.jsonl> [top_n]}"
TOP_N="${2:-5}"

if [[ ! -f "$MANIFEST" ]]; then
    echo "manifest not found: $MANIFEST" >&2
    exit 1
fi

# Manifest paths can be absolute or relative to manifest dir.
MANIFEST_DIR="$(cd "$(dirname "$MANIFEST")" && pwd)"

resolve_wav() {
    local wav="$1"
    if [[ "$wav" = /* ]]; then echo "$wav"; else echo "$MANIFEST_DIR/$wav"; fi
}

# Score one clip:
#   duration in [10.0, 13.5] = +1.0 (sweet spot, decreases linearly outside)
#   RMS dB in [-25, -15]    = +0.5 (clean speech band)
#   peak (linear, [-9, 0] dBFS): +0.5 at peak ≤ -9, ramping to 0 at peak ≥ 0
#     (CSM tracks input amplitude, so hotter context → hotter / clipping output;
#     2026-04-29 A/B showed binary peak threshold misranks hot clips against
#     each other — linear penalty differentiates them)
# Anything outside window: 0 contribution. Total max = 2.0.
score_clip() {
    local dur="$1" rms="$2" peak="$3"
    local s_dur s_rms s_peak
    if (( $(echo "$dur >= 10 && $dur <= 13.5" | bc -l) )); then
        s_dur=1.0
    elif (( $(echo "$dur >= 8 && $dur < 10" | bc -l) )); then
        s_dur=$(echo "scale=3; ($dur - 8) / 2" | bc -l)
    elif (( $(echo "$dur > 13.5 && $dur <= 16" | bc -l) )); then
        s_dur=$(echo "scale=3; (16 - $dur) / 2.5" | bc -l)
    else
        s_dur=0
    fi

    if (( $(echo "$rms >= -25 && $rms <= -15" | bc -l) )); then s_rms=0.5; else s_rms=0; fi

    # Linear peak penalty: 0.5 at -9 dBFS or quieter, 0.0 at 0 dBFS, clamped.
    if (( $(echo "$peak <= -9" | bc -l) )); then
        s_peak=0.5
    elif (( $(echo "$peak >= 0" | bc -l) )); then
        s_peak=0
    else
        s_peak=$(echo "scale=4; (- $peak) / 18" | bc -l)
    fi

    echo "scale=4; $s_dur + $s_rms + $s_peak" | bc -l
}

# Decode int16 max → dBFS roughly: int16_max = 32767 = 0 dBFS.
# Peak in dB = 20 * log10(|peak_int16| / 32767).
peak_to_db() {
    local p="$1"
    local abs_p
    abs_p=$(echo "$p" | awk '{print ($1 < 0) ? -$1 : $1}')
    if (( $(echo "$abs_p < 1" | bc -l) )); then echo "-100"; return; fi
    echo "scale=2; l($abs_p / 32767) * 20 / l(10)" | bc -l
}

# Collect (group_key, score, dur, rms, peak, wav, transcript) per row.
TMPFILE="$(mktemp)"
trap "rm -f '$TMPFILE'" EXIT

while IFS= read -r line; do
    [[ -z "$line" ]] && continue
    [[ "$line" =~ ^# ]] && continue
    wav=$(echo "$line" | jq -r '.wav')
    txt=$(echo "$line" | jq -r '.transcript')
    spk=$(echo "$line" | jq -r '.speaker // 0')
    abs_wav="$(resolve_wav "$wav")"
    [[ ! -f "$abs_wav" ]] && continue

    # Use the manifest filename (stem) as part of the group, since speaker
    # IDs are per-file.
    src_stem="$(basename "$MANIFEST" .jsonl)"
    group="${src_stem}/spk${spk}"

    dur=$(ffprobe -v error -show_entries format=duration -of csv=p=0 "$abs_wav" 2>/dev/null)
    [[ -z "$dur" ]] && continue

    # ffmpeg astats output: "Max level: <int16>" + "RMS level dB: <float>"
    stats=$(ffmpeg -i "$abs_wav" -af "astats=metadata=1:reset=0" -f null - 2>&1 || true)
    rms=$(echo "$stats" | grep -m1 "RMS level dB:" | awk -F'RMS level dB:' '{print $2}' | tr -d ' ')
    peak_int=$(echo "$stats" | grep -m1 "Max level:" | awk -F'Max level:' '{print $2}' | tr -d ' ')
    [[ -z "$rms" || -z "$peak_int" ]] && continue
    # ffmpeg might report Inf for silent files
    [[ "$rms" =~ "inf" ]] && continue
    peak_db=$(peak_to_db "$peak_int")

    score=$(score_clip "$dur" "$rms" "$peak_db")
    printf "%s\t%s\t%.2f\t%.2f\t%.2f\t%s\t%s\n" \
        "$group" "$score" "$dur" "$rms" "$peak_db" "$abs_wav" "$txt" >> "$TMPFILE"
done < "$MANIFEST"

# Group, sort within each group by score DESC, keep top N.
echo "# group                score  dur(s)  rms(dB)  peak(dB)  wav  ←  transcript"
sort -t $'\t' -k1,1 -k2,2gr "$TMPFILE" | awk -F'\t' -v n="$TOP_N" '
{
    if ($1 != prev) { count = 0; prev = $1 }
    if (count < n) {
        printf "%s\t%s\t%s\t%s\t%s\t%s\t←  %s\n", $1, $2, $3, $4, $5, $6, $7
        count++
    }
}'