#!/bin/bash
# pepito transcription command, example based on WhisperKit (whisperkit-cli),
# with speaker labels ([A], [B], ...) from its speaker diarization.
# Requires pepito 26.61 or later, and a Mac with Apple silicon.
#
# Set the path of this file as the transcription command in pepito (or a
# command line that runs it). pepito passes the job in environment
# variables (PEPITO_INPUT_PATH, PEPITO_OUTPUT_PATH, ...).
# It must write an SRT file to PEPITO_OUTPUT_PATH and exit 0.
# Any non-zero exit code is reported by pepito as a failed transcription,
# with the end of stderr as the message.
#
# This script never downloads anything: install whisperkit-cli and download
# the models beforehand (see the page that publishes this file).
#
# Not used here: PEPITO_TASK and the other variables of the recording
# (absent when Settings runs the Test button).
set -euo pipefail

# Settings specific to this script (outside the PEPITO_ namespace, which
# pepito reserves and clears); pepito ignores them.
MODEL="${WHISPERKIT_MODEL:-large-v3-v20240930_turbo_632MB}"
MODELS_DIR="${WHISPERKIT_MODELS_DIR:-$HOME/Library/Application Support/whisperkit-models}"
NUM_SPEAKERS="${WHISPERKIT_NUM_SPEAKERS:-}"

WHISPER_MODEL_PATH="$MODELS_DIR/models/argmaxinc/whisperkit-coreml/openai_whisper-$MODEL"
DIARIZATION_MODEL_PATH="$MODELS_DIR/models/argmaxinc/speakerkit-coreml"

# Preflight checks. whisperkit-cli downloads whatever is missing, the
# tokenizer included, so the script refuses to start without them.
if ! command -v whisperkit-cli >/dev/null 2>&1; then
  echo "whisperkit-cli not found in PATH (brew install whisperkit-cli)" >&2
  exit 1
fi
if [ ! -d "$WHISPER_MODEL_PATH" ]; then
  echo "model not found: $WHISPER_MODEL_PATH" >&2
  exit 1
fi
if [ ! -d "$DIARIZATION_MODEL_PATH" ]; then
  echo "diarization models not found: $DIARIZATION_MODEL_PATH" >&2
  exit 1
fi
if ! ls "$MODELS_DIR"/models/openai/*/tokenizer.json >/dev/null 2>&1; then
  echo "tokenizer not found in: $MODELS_DIR/models/openai" >&2
  exit 1
fi
case "${PEPITO_OUTPUT_PATH:-}" in
  *.srt) ;;
  *)
    echo "PEPITO_OUTPUT_PATH must end with .srt: ${PEPITO_OUTPUT_PATH:-<unset>}" >&2
    exit 1
    ;;
esac

# The current directory is pepito's working folder, deleted at the end.
# whisperkit-cli prints its progress and the transcript on stdout, hence
# >/dev/null (pepito logs stdout privately, but there is no reason to print
# the transcript). Errors stay on stderr.

# 1. Speaker turns, as RTTM.
diarize=(whisperkit-cli diarize --audio-path "$PEPITO_INPUT_PATH"
  --rttm-path "$PWD/diarization.rttm" --model-path "$DIARIZATION_MODEL_PATH")
if [ -n "$NUM_SPEAKERS" ]; then
  diarize+=(--num-speakers "$NUM_SPEAKERS")
fi
"${diarize[@]}" >/dev/null

# 2. Transcript with word timestamps: --report writes an SRT with one
# segment per word in --report-path.
# Language: PEPITO_LANGUAGE_CODE is an ISO 639-1 code ("fr") or "auto";
# without --language, Whisper detects the language.
transcribe=(whisperkit-cli transcribe --audio-path "$PEPITO_INPUT_PATH"
  --model-path "$WHISPER_MODEL_PATH" --download-tokenizer-path "$MODELS_DIR"
  --word-timestamps --skip-special-tokens --report --report-path "$PWD/report")
if [ "$PEPITO_LANGUAGE_CODE" != "auto" ]; then
  transcribe+=(--language "$PEPITO_LANGUAGE_CODE")
fi
mkdir -p "$PWD/report"
"${transcribe[@]}" >/dev/null

words=("$PWD"/report/*.srt)
if [ ! -f "${words[0]}" ]; then
  echo "whisperkit-cli wrote no SRT in $PWD/report" >&2
  exit 1
fi

# 3. Each word goes to the speaker whose turn overlaps it most (the nearest
# turn if none does). A new segment starts on a change of speaker, after a
# sentence end (. ? !) or after a silence of more than 1 s. A token without
# a leading space continues the previous word: same speaker, no cut.
# Labels are added only when more than one speaker is found.
LC_ALL=C awk -v GAP=1.0 '
function secs(t,  p) { split(t, p, /[:,]/); return p[1] * 3600 + p[2] * 60 + p[3] + p[4] / 1000 }
function stamp(t,  ms) {
  ms = int(t * 1000 + 0.5)
  return sprintf("%02d:%02d:%02d,%03d", int(ms / 3600000), int(ms / 60000) % 60, int(ms / 1000) % 60, ms % 1000)
}
function speaker(s, e,  i, o, d, best, bo, bd, mid) {
  bo = 0; best = ""
  for (i = 1; i <= nt; i++) {
    o = (e < te[i] ? e : te[i]) - (s > ts[i] ? s : ts[i])
    if (o > bo) { bo = o; best = tk[i] }
  }
  if (best != "" || nt == 0) return best
  mid = (s + e) / 2; bd = -1
  for (i = 1; i <= nt; i++) {
    d = mid - ts[i]; if (d < 0) d = -d
    o = mid - te[i]; if (o < 0) o = -o
    if (o < d) d = o
    if (bd < 0 || d < bd) { bd = d; best = tk[i] }
  }
  return best
}
function word(s, e, t,  bare, cont, k) {
  bare = t; gsub(/^[ \t]+|[ \t]+$/, "", bare)
  if (bare == "" || bare == "...") return
  cont = n > 0 && substr(t, 1, 1) != " "
  k = cont ? cs[n] : speaker(s, e)
  if (n == 0 || (!cont && (k != cs[n] || stop || s - ce[n] > GAP))) {
    n++; cb[n] = s; cs[n] = k; ct[n] = ""
    if (k != "") seen[k] = 1
  }
  ce[n] = e; ct[n] = ct[n] t
  stop = bare ~ /[.?!]["»”)]*$/
}
FNR == NR {
  if ($1 == "SPEAKER") { nt++; ts[nt] = $4 + 0; te[nt] = $4 + $5; tk[nt] = $8 }
  next
}
/-->/ { s = secs($1); e = secs($3); text = ""; intext = 1; next }
intext && /^$/ { word(s, e, text); intext = 0; next }
intext { text = text $0; next }
END {
  if (intext) word(s, e, text)
  for (k in seen) ns++
  for (i = 1; i <= n; i++) {
    t = ct[i]; gsub(/^[ \t]+|[ \t]+$/, "", t)
    if (t == "") continue
    if (ns > 1 && cs[i] != "") t = "[" cs[i] "] " t
    printf "%d\n%s --> %s\n%s\n\n", ++m, stamp(cb[i]), stamp(ce[i]), t
  }
}
' "$PWD/diarization.rttm" "${words[0]}" > "$PEPITO_OUTPUT_PATH"