#!/usr/bin/env bash
# transcribe — offline audio/video transcription via whisper.cpp
#
# Usage:
#   transcribe [options] <file> [file ...]
#
# Options:
#   -m, --model <name>     Model to use (default: large-v3-turbo-q5_0)
#                          Examples: tiny, base, small, medium, large-v3,
#                                    large-v3-turbo, large-v3-turbo-q5_0
#   -l, --language <code>  Language: en, hu, auto (default: auto)
#   -f, --format <fmt>     Output format: txt|srt|vtt|json|all (default: txt,srt)
#   -o, --output-dir <dir> Where to put outputs (default: alongside input)
#   -t, --threads <n>      CPU threads (default: physical core count)
#       --translate        Translate to English
#       --diarize          Stereo diarization (2-channel only)
#       --list-models      List installed models and exit
#       --download <name>  Force-download a model and exit
#   -h, --help             Show this help

set -euo pipefail

# shellcheck disable=SC1090
source "${HOME}/.local/share/chezmoi/scripts/utils" 2>/dev/null || {
    cmd_exists() { command -v "$1" >/dev/null 2>&1; }
}

MODEL_DIR="${WHISPER_MODEL_DIR:-$HOME/.local/share/whisper-models}"
DEFAULT_MODEL="large-v3-turbo-q5_0"
HF_BASE="https://huggingface.co/ggerganov/whisper.cpp/resolve/main"

die() {
    printf '\033[31merror:\033[0m %s\n' "$*" >&2
    exit 1
}
info() { printf '\033[36m==>\033[0m %s\n' "$*" >&2; }

# whisper.cpp installs as `whisper-cli` (newer) or `main` (older); detect.
detect_whisper_bin() {
    for b in whisper-cli whisper-cpp main; do
        if command -v "$b" >/dev/null 2>&1; then
            echo "$b"
            return
        fi
    done
    die "whisper.cpp not found. Install with: brew install whisper-cpp"
}

ensure_model() {
    local model="$1"
    local file="$MODEL_DIR/ggml-${model}.bin"
    mkdir -p "$MODEL_DIR"
    if [[ -f "$file" ]]; then
        echo "$file"
        return
    fi
    info "Downloading model: ggml-${model}.bin (this happens once)"
    if cmd_exists gum; then
        gum spin --spinner dot --title "Fetching ${model}…" -- \
            curl -L --fail -o "$file.tmp" "$HF_BASE/ggml-${model}.bin"
    else
        curl -L --fail --progress-bar -o "$file.tmp" "$HF_BASE/ggml-${model}.bin"
    fi
    mv "$file.tmp" "$file"
    echo "$file"
}

list_models() {
    info "Models in $MODEL_DIR:"
    if [[ -d "$MODEL_DIR" ]]; then
        ls -lh "$MODEL_DIR"/ggml-*.bin 2>/dev/null | awk '{print "  " $9 " (" $5 ")"}'
    else
        echo "  (none)"
    fi
}

# Defaults
MODEL="$DEFAULT_MODEL"
LANG="auto"
FORMAT="txt,srt"
OUT_DIR=""
THREADS="$(sysctl -n hw.physicalcpu 2>/dev/null || echo 8)"
EXTRA_ARGS=()
FILES=()

while [[ $# -gt 0 ]]; do
    case "$1" in
    -m | --model)
        MODEL="$2"
        shift 2
        ;;
    -l | --language)
        LANG="$2"
        shift 2
        ;;
    -f | --format)
        FORMAT="$2"
        shift 2
        ;;
    -o | --output-dir)
        OUT_DIR="$2"
        shift 2
        ;;
    -t | --threads)
        THREADS="$2"
        shift 2
        ;;
    --translate)
        EXTRA_ARGS+=(--translate)
        shift
        ;;
    --diarize)
        EXTRA_ARGS+=(--diarize)
        shift
        ;;
    --list-models)
        list_models
        exit 0
        ;;
    --download)
        ensure_model "$2" >/dev/null
        info "OK"
        exit 0
        ;;
    -h | --help)
        sed -n '2,/^$/p' "$0" | sed 's/^# \?//'
        exit 0
        ;;
    --)
        shift
        FILES+=("$@")
        break
        ;;
    -*) die "Unknown option: $1" ;;
    *)
        FILES+=("$1")
        shift
        ;;
    esac
done

[[ ${#FILES[@]} -eq 0 ]] && die "No input files. See: transcribe --help"
cmd_exists ffmpeg || die "ffmpeg is required (brew install ffmpeg)"

WHISPER_BIN="$(detect_whisper_bin)"
MODEL_FILE="$(ensure_model "$MODEL")"

# Build whisper output flags
declare -a OUT_FLAGS=()
IFS=',' read -ra FMTS <<<"$FORMAT"
for f in "${FMTS[@]}"; do
    case "$f" in
    txt) OUT_FLAGS+=(--output-txt) ;;
    srt) OUT_FLAGS+=(--output-srt) ;;
    vtt) OUT_FLAGS+=(--output-vtt) ;;
    json) OUT_FLAGS+=(--output-json) ;;
    all) OUT_FLAGS+=(--output-txt --output-srt --output-vtt --output-json) ;;
    *) die "Unknown format: $f" ;;
    esac
done

TMPDIR_WORK="$(mktemp -d -t transcribe.XXXXXX)"
trap 'rm -rf "$TMPDIR_WORK"' EXIT

for input in "${FILES[@]}"; do
    [[ -f "$input" ]] || {
        echo "skip (not a file): $input" >&2
        continue
    }

    base="$(basename "${input%.*}")"
    out_base="${OUT_DIR:+$OUT_DIR/}${OUT_DIR:-$(dirname "$input")/}${base}"
    # If OUT_DIR set, ensure it exists
    [[ -n "$OUT_DIR" ]] && mkdir -p "$OUT_DIR"

    wav="$TMPDIR_WORK/${base}.wav"

    info "Preparing: $input → 16kHz mono WAV"
    ffmpeg -hide_banner -loglevel error -y -i "$input" \
        -ac 1 -ar 16000 -c:a pcm_s16le "$wav"

    info "Transcribing with $WHISPER_BIN ($MODEL, lang=$LANG, threads=$THREADS)"
    "$WHISPER_BIN" \
        --model "$MODEL_FILE" \
        --language "$LANG" \
        --threads "$THREADS" \
        --output-file "${OUT_DIR:+$OUT_DIR/}${base}" \
        "${OUT_FLAGS[@]}" \
        "${EXTRA_ARGS[@]}" \
        --file "$wav"

    info "Done: ${out_base}.{${FORMAT}}"
done