#!/usr/bin/env bash

set -euo pipefail
umask 077

srt=false
diarization=true
chunk_minutes=
jobs=8
speaker_model=
speakers=
diagnostics=
usage() {
    echo "Usage: ${0##*/} [--srt] [--no-diarization] [--speakers N] [--speaker-model MODEL] [--chunk-minutes 1..30] [--jobs 1..64] [--diagnostics DIR] AUDIO_FILE"
    echo 'Transcribe audio with automatic speaker detection; print wrapped text to stdout.'
    echo 'MAI transcribes through OpenRouter; Pyannote labels speakers across the full recording.'
    echo '  --srt  Print SRT subtitles in short cues using word timestamps instead.'
    echo '  --no-diarization  Skip Pyannote and print text without speaker labels.'
    echo '  --speakers N  Optional known speaker count (1..64; otherwise detected automatically).'
    echo '  --speaker-model MODEL  Hosted Pyannote model: precision-2 (default) or community-1.'
    echo '  --chunk-minutes N  Maximum transcription upload duration in minutes (default: 30).'
    echo '  --jobs N  Maximum chunks processed in parallel (default: 8; 1 for sequential).'
    echo '  --diagnostics DIR  Save timing and speaker JSON in a new private directory.'
}
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:-} != precision-2 && ${2:-} != community-1 ]]; then
                echo '--speaker-model must be precision-2 or community-1' >&2
                exit 1
            fi
            speaker_model=$2
            shift 2
            ;;
        --diagnostics)
            if [[ -z ${2:-} || $2 == -* ]]; then
                echo '--diagnostics requires a new directory path' >&2
                exit 1
            fi
            diagnostics=$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 $speaker_model || -n $speakers ]] && ! "$diarization"; then
    echo '--speakers and --speaker-model cannot be combined with --no-diarization' >&2
    exit 1
fi
if "$diarization" && [[ -z $speaker_model ]]; then
    speaker_model=precision-2
fi
if [[ -z $chunk_minutes ]]; then
    # Speaker identities come from a separate whole-recording timeline.
    chunk_minutes=30
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
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

PYANNOTE_API_KEY_REF="op://API Keys/pyannote-api/credential"

# Explicit environment keys also work in headless sessions without 1Password.
# Resolve both credentials before starting any paid requests.
if "$diarization" && [[ -z ${PYANNOTE_API_KEY:-} ]]; then
    if [[ -n ${PYANNOTE_API_KEY_REF:-} ]] && command -v op >/dev/null; then
        PYANNOTE_API_KEY=$(op read "$PYANNOTE_API_KEY_REF")
    else
        echo 'Set PYANNOTE_API_KEY or PYANNOTE_API_KEY_REF for Pyannote speaker diarization.' >&2
        exit 1
    fi
fi
if [[ -z ${OPENROUTER_API_KEY:-} ]]; then
    if ! command -v op >/dev/null; then
        echo 'Set OPENROUTER_API_KEY, or install and sign in to the 1Password CLI.' >&2
        exit 1
    fi
    OPENROUTER_API_KEY=$(op read "${OPENROUTER_API_KEY_REF:-op://API Keys/openrouter-mai-transcribe-2/credential}")
fi
if [[ ! $OPENROUTER_API_KEY =~ [^[:space:]] || $OPENROUTER_API_KEY == *$'\n'* || $OPENROUTER_API_KEY == *$'\r'* ]]; then
    echo 'OPENROUTER_API_KEY must be a nonempty single-line API key.' >&2
    exit 1
fi
if "$diarization"; then
    if [[ ! $PYANNOTE_API_KEY =~ [^[:space:]] || $PYANNOTE_API_KEY == *$'\n'* || $PYANNOTE_API_KEY == *$'\r'* ]]; then
        echo 'PYANNOTE_API_KEY must be a nonempty single-line API key.' >&2
        exit 1
    fi
    export PYANNOTE_API_KEY
fi
export OPENROUTER_API_KEY
if [[ -n $diagnostics ]]; then
    mkdir -m 700 -- "$diagnostics"
    diagnostics=$(realpath -- "$diagnostics")
fi

save_diagnostic() {
    if [[ -n $diagnostics ]]; then
        cp -- "$1" "$diagnostics/$2"
    fi
}

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")")
printf 'Authorization: Bearer %s\n' "$OPENROUTER_API_KEY" >"$tmp/openrouter.header"

# 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 -B "$script_dir/chunks.py" plan "$tmp/audio.wav" "$chunk_minutes" >"$tmp/chunks.json"
save_diagnostic "$tmp/chunks.json" 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" \
        '{model: "microsoft/mai-transcribe-2", input_audio: {data: $data, format: "mp3"},
          provider: {options: {azure: {diarization: {enabled: false}}}}}
         + (if $timestamps then {response_format: "verbose_json", timestamp_granularities: ["segment", "word"]}
            else {} end)' \
        >"$work/request.json"

    curl --silent --show-error --fail-with-body --connect-timeout 30 --max-time 1800 \
        https://openrouter.ai/api/v1/audio/transcriptions \
        -H "@$tmp/openrouter.header" \
        -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
    # Ignore unsolicited request-local labels, including in the explicit opt-out.
    jq 'del(.words[]?.speaker, .segments[]?.speaker)' "$tmp/response-$i.json" >"$work/unlabelled.json"
    mv -- "$work/unlabelled.json" "$tmp/response-$i.json"
    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 "$diarization" || "$srt" || ((count > 1)); then
    python3 -B "$script_dir/chunks.py" merge "$tmp/chunks.json" --json >"$tmp/transcript.json"
else
    cp -- "$tmp/response-0.json" "$tmp/transcript.json"
fi
save_diagnostic "$tmp/transcript.json" transcript.json
if "$diarization"; then
    speaker_options=()
    if [[ -n $speakers ]]; then speaker_options+=(--speakers "$speakers"); fi
    label_speakers() {
        trap - EXIT INT TERM
        python3 -B "$script_dir/diarization.py" run "$tmp/audio.wav" "$tmp/voices" \
            --model "$speaker_model" "${speaker_options[@]}"
        save_diagnostic "$tmp/voices/result.json" diarization.json
        python3 -B "$script_dir/diarization.py" apply "$tmp/transcript.json" \
            "$tmp/voices/result.json" --json "${speaker_options[@]}" >"$tmp/labeled.json"
        save_diagnostic "$tmp/labeled.json" labeled.json
    }
    # The encoder and HTTP client get the same interruption cleanup as STT.
    label_speakers &
    workers[$!]=speaker-diarization
    wait_for_worker
    rendered_source=$tmp/labeled.json
else
    rendered_source=$tmp/transcript.json
fi
if "$srt"; then
    python3 -B "$script_dir/subtitles.py" "$rendered_source" >"$tmp/output"
else
    python3 -B "$script_dir/transcript.py" "$rendered_source" >"$tmp/output"
fi
cat -- "$tmp/output"
set +m
