#!/usr/bin/env bash
# Transcribes a video/audio file using whisper.cpp in a podman container, using
# hardware acceleration when it's available on the host:
#   - NVIDIA GPU present  -> ghcr.io/ggml-org/whisper.cpp:main-cuda   (CDI nvidia device)
#   - /dev/dri + vainfo    -> ghcr.io/ggml-org/whisper.cpp:main-vulkan (AMD/Intel, Vulkan)
#   - neither               -> ghcr.io/ggml-org/whisper.cpp:main       (CPU)
#
# If the accelerated container fails to run (missing CDI setup, no such image tag on
# this whisper.cpp release, etc.) this script automatically falls back to the CPU
# image rather than failing outright, and says so on stderr.
#
# Usage: transcribe.sh <media-file> <output-dir> [model]
#   model defaults to "base.en" (small, fast, English). Pass e.g. "small" for a
#   multilingual model, or "medium.en" for higher accuracy.
#   On a multi-GPU AMD/Intel host, set WHISPER_VULKAN_DEVICE to the desired render
#   node, e.g. WHISPER_VULKAN_DEVICE=/dev/dri/renderD128.
#
# Writes:
#   <output-dir>/<basename>.transcript.txt
#   <output-dir>/<basename>.transcript.srt

set -euo pipefail

usage() {
  echo "Usage: $(basename "$0") <media-file> <output-dir> [model=base.en]" >&2
  exit 1
}

[ $# -ge 2 ] || usage

MEDIA="$1"
OUT_DIR="$2"
MODEL="${3:-base.en}"

[ -f "$MEDIA" ] || { echo "No such file: $MEDIA" >&2; exit 1; }
command -v podman >/dev/null 2>&1 || { echo "podman is required for transcription" >&2; exit 1; }
command -v curl >/dev/null 2>&1 || { echo "curl is required to fetch the whisper model" >&2; exit 1; }

MEDIA="$(cd "$(dirname "$MEDIA")" && pwd)/$(basename "$MEDIA")"
BASENAME="$(basename "${MEDIA%.*}")"
WORKDIR="$(mktemp -d)"
trap 'rm -rf "$WORKDIR"' EXIT

WAV="$WORKDIR/audio.wav"

echo "Extracting audio..." >&2
if command -v ffmpeg >/dev/null 2>&1; then
  ffmpeg -y -i "$MEDIA" -ar 16000 -ac 1 -c:a pcm_s16le "$WAV" -loglevel error
else
  echo "ffmpeg not found on host, extracting via podman instead" >&2
  podman run --rm \
    -v "$(dirname "$MEDIA")":/in:Z \
    -v "$WORKDIR":/out:Z \
    jrottenberg/ffmpeg:6-alpine \
    -y -i "/in/$(basename "$MEDIA")" -ar 16000 -ac 1 -c:a pcm_s16le /out/audio.wav -loglevel error
fi

MODEL_DIR="${XDG_CACHE_HOME:-$HOME/.cache}/ticket-skill/whisper-models"
mkdir -p "$MODEL_DIR"
MODEL_FILE="$MODEL_DIR/ggml-${MODEL}.bin"
if [ ! -f "$MODEL_FILE" ]; then
  echo "Downloading whisper model '${MODEL}' (first run only, cached under ${MODEL_DIR})..." >&2
  curl -sSL -o "$MODEL_FILE" "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-${MODEL}.bin"
fi

ACCEL="cpu"
IMAGE="ghcr.io/ggml-org/whisper.cpp:main"
DEVICE_ARGS=()

if command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L >/dev/null 2>&1; then
  ACCEL="cuda"
  IMAGE="ghcr.io/ggml-org/whisper.cpp:main-cuda"
  DEVICE_ARGS=(--device nvidia.com/gpu=all)
elif [ -d /dev/dri ] && command -v vainfo >/dev/null 2>&1; then
  ACCEL="vulkan"
  IMAGE="ghcr.io/ggml-org/whisper.cpp:main-vulkan"
  if [ -n "${WHISPER_VULKAN_DEVICE:-}" ]; then
    VULKAN_DEVICE="$(readlink -f "$WHISPER_VULKAN_DEVICE")"
    [ -e "$VULKAN_DEVICE" ] || {
      echo "WHISPER_VULKAN_DEVICE does not exist: $WHISPER_VULKAN_DEVICE" >&2
      exit 1
    }
    case "$VULKAN_DEVICE" in
      /dev/dri/renderD*) ;;
      *)
        echo "WHISPER_VULKAN_DEVICE must resolve to /dev/dri/renderD*: $VULKAN_DEVICE" >&2
        exit 1
        ;;
    esac
    DEVICE_ARGS=(--device "$VULKAN_DEVICE")
  else
    DEVICE_ARGS=(--device /dev/dri)
  fi
fi

echo "Selected acceleration: ${ACCEL} (image: ${IMAGE})" >&2

run_whisper() {
  local image="$1"
  shift
  podman run --rm "$@" \
    --entrypoint /app/build/bin/whisper-cli \
    -v "$WORKDIR":/data:Z \
    -v "$MODEL_DIR":/models:Z \
    "$image" \
    -m "/models/ggml-${MODEL}.bin" -f /data/audio.wav -otxt -osrt -of /data/transcript
}

SUCCESS=0
if [ "$ACCEL" != "cpu" ]; then
  if run_whisper "$IMAGE" "${DEVICE_ARGS[@]}"; then
    SUCCESS=1
  else
    echo "Accelerated transcription (${ACCEL}) failed, falling back to CPU." >&2
  fi
fi

if [ "$SUCCESS" -eq 0 ]; then
  run_whisper "ghcr.io/ggml-org/whisper.cpp:main"
fi

[ -f "$WORKDIR/transcript.txt" ] || { echo "Transcription did not produce output" >&2; exit 1; }

mkdir -p "$OUT_DIR"
cp "$WORKDIR/transcript.txt" "$OUT_DIR/${BASENAME}.transcript.txt"
[ -f "$WORKDIR/transcript.srt" ] && cp "$WORKDIR/transcript.srt" "$OUT_DIR/${BASENAME}.transcript.srt"

echo "Wrote ${OUT_DIR}/${BASENAME}.transcript.txt"