#!/usr/bin/env bash

set -euo pipefail

srt=false
diarization=true
chunk_minutes=
jobs=8
speaker_model=
speakers=
usage() {
    echo "Usage: ${0##*/} [--srt] [--no-diarization] [--speakers N] [--speaker-model MODEL] [--chunk-minutes 1..30] [--jobs 1..64] AUDIO_FILE"
    echo 'Transcribe audio using OPENROUTER_API_KEY; print wrapped text with speaker markers to stdout.'
    echo '  --srt  Print SRT subtitles in short cues using word timestamps instead.'
    echo '  --no-diarization  Disable speaker diarization (enabled by default).'
    echo '  --speakers N  Reconcile voices through OpenRouter with a known count (1..64).'
    echo '  --speaker-model MODEL  Reconcile voices with an audio chat model (count optional).'
    echo '    --speakers defaults to google/gemini-3.8-flash; adds one audio request.'
    echo '  --chunk-minutes N  Maximum upload duration in minutes (default: 10; 30 with --no-diarization).'
    echo '  --jobs N  Maximum chunks processed in parallel (default: 8; 1 for sequential).'
}
while [[ $# -gt 0 ]]; do
    case $1 in
        --srt) srt=true; shift ;;
        --no-diarization) diarization=false; shift ;;
        --speakers)
            if [[ ! ${2:-} =~ ^([1-9]|[1-5][0-9]|6[0-4])$ ]]; then
                echo '--speakers must be an integer between 1 and 64' >&2
                exit 1
            fi
            speakers=$2
            shift 2
            ;;
        --speaker-model)
            if [[ ! ${2:-} =~ ^~?[[:alnum:]][[:alnum:]._:/-]*$ ]]; then
                echo '--speaker-model requires an OpenRouter audio chat model ID' >&2
                exit 1
            fi
            speaker_model=$2
            shift 2
            ;;
        --jobs)
            if [[ ! ${2:-} =~ ^([1-9]|[1-5][0-9]|6[0-4])$ ]]; then
                echo '--jobs must be an integer between 1 and 64' >&2
                exit 1
            fi
            jobs=$2
            shift 2
            ;;
        --chunk-minutes)
            if [[ ! ${2:-} =~ ^([1-9]|[12][0-9]|30)$ ]]; then
                echo '--chunk-minutes must be an integer between 1 and 30' >&2
                exit 1
            fi
            chunk_minutes=$2
            shift 2
            ;;
        -h|--help) usage; exit 0 ;;
        --) shift; break ;;
        -*) usage >&2; exit 1 ;;
        *) break ;;
    esac
done
if [[ -n $speakers && -z $speaker_model ]]; then
    speaker_model=google/gemini-3.8-flash
fi
if [[ -n $speaker_model ]] && ! "$diarization"; then
    echo '--speakers and --speaker-model cannot be combined with --no-diarization' >&2
    exit 1
fi
if [[ -z $chunk_minutes ]]; then
    # Azure documents diarization failures for recordings around 15 minutes.
    if "$diarization"; then chunk_minutes=10; else chunk_minutes=30; fi
fi
if [[ $# != 1 ]]; then
    usage >&2
    exit 1
fi
if [[ ! -f $1 || ! -r $1 ]]; then
    echo "Cannot read audio file: $1" >&2
    exit 1
fi
if (( BASH_VERSINFO[0] < 5 || (BASH_VERSINFO[0] == 5 && BASH_VERSINFO[1] < 1) )); then
    echo 'Bash 5.1 or newer is required for parallel chunk processing.' >&2
    exit 1
fi
export OPENROUTER_API_KEY
OPENROUTER_API_KEY="$(op read "op://API Keys/openrouter-mai-transcribe-2/credential")"
: "${OPENROUTER_API_KEY:?Set OPENROUTER_API_KEY first}"
dependencies=(ffmpeg curl jq base64 mktemp python3)
for tool in "${dependencies[@]}"; do
    command -v "$tool" >/dev/null || {
        echo "Missing dependency: $tool" >&2
        exit 1
    }
done

tmp=$(mktemp -d)
declare -A workers=()
cleanup() {
    trap - EXIT
    trap '' INT TERM
    local pid
    local -a pending
    # Include a job interrupted between launch and registration in workers.
    mapfile -t pending < <(jobs -p)
    # Each worker owns a process group, including its encoder and curl. Stop
    # the whole group and reap workers before removing files they might use.
    for pid in "${pending[@]}"; do
        kill -KILL -- "-$pid" 2>/dev/null || true
    done
    for pid in "${pending[@]}"; do
        wait "$pid" 2>/dev/null || true
    done
    rm -rf -- "$tmp"
}
trap cleanup EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
script_dir=$(dirname -- "$(realpath -- "$0")")

# Decode once so duration, pause detection, and all cuts share a sample clock.
# This also works for inputs whose container omits or misreports the duration.
ffmpeg -nostdin -hide_banner -loglevel error -i "$(realpath -- "$1")" \
    -map 0:a:0 -ac 1 -ar 16000 -c:a pcm_s16le "$tmp/audio.wav"
python3 "$script_dir/chunks.py" plan "$tmp/audio.wav" "$chunk_minutes" >"$tmp/chunks.json"
count=$(jq 'length' "$tmp/chunks.json")
timestamps=$srt
if "$diarization"; then timestamps=true; fi
if (( count > 1 )); then
    timestamps=true
    echo "Splitting audio into $count overlapping chunks (maximum $chunk_minutes minutes each)." >&2
    echo "Processing up to $jobs chunks in parallel." >&2
fi

transcribe_chunk() {
    trap - EXIT INT TERM
    local i=$1 start duration work
    work=$tmp/chunk-$i
    mkdir -- "$work"
    read -r start duration < <(jq -r --argjson i "$i" \
        '.[$i] | "\(.start) \(.end - .start)"' "$tmp/chunks.json")
    if (( count > 1 )); then
        echo "Transcribing chunk $((i + 1))/$count (start ${start}s, duration ${duration}s)..." >&2
    fi
    ffmpeg -nostdin -hide_banner -loglevel error -y -ss "$start" -i "$tmp/audio.wav" \
        -t "$duration" -map 0:a:0 -ac 1 -ar 16000 -c:a libmp3lame -b:a 64k "$work/audio.mp3"

    # Keep large audio payloads out of shell arguments (which have a size limit).
    base64 <"$work/audio.mp3" | tr -d '\n\r' >"$work/audio.b64"
    jq -n --rawfile data "$work/audio.b64" --argjson timestamps "$timestamps" \
        --argjson diarization "$diarization" \
        '{model: "microsoft/mai-transcribe-2", input_audio: {data: $data, format: "mp3"},
          provider: {options: {azure: {diarization: {enabled: $diarization}}}}}
         + (if $timestamps then {response_format: "verbose_json", timestamp_granularities: ["segment", "word"]}
            else {} end)' \
        >"$work/request.json"

    curl --silent --show-error --fail-with-body \
        https://openrouter.ai/api/v1/audio/transcriptions \
        -H "Authorization: Bearer $OPENROUTER_API_KEY" \
        -H 'Content-Type: application/json' \
        --data-binary @"$work/request.json" -o "$tmp/response-$i.json" || {
        echo "Transcription failed for chunk $((i + 1))/$count." >&2
        if [[ -f $tmp/response-$i.json ]]; then
            cat "$tmp/response-$i.json" >&2
        fi
        exit 1
    }
    jq -e 'if .error then error("Transcription failed: " + (.error | tostring))
        elif (.text | type) == "string" then true
        else error("Transcription failed: missing text in response") end' \
        "$tmp/response-$i.json" >/dev/null
    if "$diarization" && jq -e '.text | test("\\S")' "$tmp/response-$i.json" >/dev/null; then
        if ! jq -e 'any(.words[]?; .speaker != null)' "$tmp/response-$i.json" >/dev/null; then
            echo "Warning: chunk $((i + 1))/$count returned no speaker labels despite diarization being enabled." >&2
        elif jq -e 'any(.words[]?; .speaker == null)' "$tmp/response-$i.json" >/dev/null; then
            echo "Warning: chunk $((i + 1))/$count returned some words without speaker labels." >&2
        fi
    fi
    rm -rf -- "$work"
}

wait_for_worker() {
    local pid finished status=0
    # wait -n skips jobs that finished before it was called. Reap those by PID
    # first, so even an immediate encoder/request failure cannot disappear.
    for pid in "${!workers[@]}"; do
        if ! kill -0 "$pid" 2>/dev/null; then
            wait "$pid" || status=$?
            unset 'workers[$pid]'
            return "$status"
        fi
    done
    wait -n -f -p finished || status=$?
    if [[ -z ${finished:-} ]]; then
        # All remaining jobs may have exited between the probe and wait -n.
        for pid in "${!workers[@]}"; do
            status=0
            wait "$pid" || status=$?
            finished=$pid
            break
        done
    fi
    unset 'workers[$finished]'
    return "$status"
}

# Job control gives every background worker a separate process group. Their
# response filenames retain recording order regardless of completion order.
set -m
next=0
while (( next < count || ${#workers[@]} > 0 )); do
    while (( next < count && ${#workers[@]} < jobs )); do
        transcribe_chunk "$next" &
        workers[$!]=$next
        next=$((next + 1))
    done
    if wait_for_worker; then
        :
    else
        status=$?
        echo 'Transcription failed; stopping remaining workers.' >&2
        exit "$status"
    fi
done
# Only emit a transcript after every chunk succeeds, including SRT conversion.
if [[ -n $speaker_model ]]; then
    # Keep request-local labels separate until the shared audio comparison.
    # The text seam matcher still removes duplicated overlap words as usual.
    python3 "$script_dir/chunks.py" merge "$tmp/chunks.json" --json --scope-speakers \
        >"$tmp/transcript.json"
    speaker_options=()
    if [[ -n $speakers ]]; then speaker_options+=(--speakers "$speakers"); fi
    reconcile_speakers() {
        trap - EXIT INT TERM
        python3 "$script_dir/diarization.py" prepare "$tmp/transcript.json" \
            "$tmp/audio.wav" "$tmp/voices" --model "$speaker_model" "${speaker_options[@]}"
        if [[ -f $tmp/voices/request.json ]]; then
            echo "Comparing speaker voices with $speaker_model through OpenRouter..." >&2
            curl --silent --show-error --fail-with-body \
                https://openrouter.ai/api/v1/chat/completions \
                -H "Authorization: Bearer $OPENROUTER_API_KEY" \
                -H 'Content-Type: application/json' \
                --data-binary @"$tmp/voices/request.json" -o "$tmp/voices/response.json" || {
                echo 'Speaker comparison failed.' >&2
                if [[ -f $tmp/voices/response.json ]]; then
                    cat "$tmp/voices/response.json" >&2
                fi
                exit 1
            }
            format=()
            if "$srt"; then format+=(--srt); fi
            python3 "$script_dir/diarization.py" apply "$tmp/transcript.json" \
                "$tmp/voices/samples.json" "$tmp/voices/response.json" \
                "${speaker_options[@]}" "${format[@]}" >"$tmp/output"
        elif "$srt"; then
            python3 "$script_dir/subtitles.py" "$tmp/transcript.json" >"$tmp/output"
        else
            python3 "$script_dir/transcript.py" "$tmp/transcript.json" >"$tmp/output"
        fi
    }
    # Give the comparison and its encoders/curl the same interruption cleanup
    # as transcription workers. Buffer rendering until this worker succeeds.
    reconcile_speakers &
    workers[$!]=speaker-comparison
    wait_for_worker
    cat -- "$tmp/output"
elif (( count > 1 )); then
    format=()
    if "$srt"; then format+=(--srt); fi
    python3 "$script_dir/chunks.py" merge "$tmp/chunks.json" "${format[@]}"
elif "$srt"; then
    python3 "$script_dir/subtitles.py" "$tmp/response-0.json"
else
    python3 "$script_dir/transcript.py" "$tmp/response-0.json"
fi
set +m
