#!/usr/bin/env python3
"""Transcribe a media file to WebVTT using faster-whisper.

Usage: transcribe.py <mediafile>

The VTT is written next to the input file as <base>.<detected-lang>.vtt
and its path is printed as a final "TRANSCRIPT_VTT:<path>" line so the
calling server can pick it up.

Environment variables:
  WHISPER_MODEL   faster-whisper model name (default: large-v3-turbo)
  WHISPER_DEVICE  "cuda" or "cpu" (default: try cuda, fall back to cpu)

Exit codes: 0 ok, 2 usage/input error, 3 faster-whisper not installed.
"""

import ctypes
import glob
import os
import sys


def preload_nvidia_libs():
    """Load pip-installed cuBLAS/cuDNN so ctranslate2 finds them without
    LD_LIBRARY_PATH (dlopen by soname succeeds once they are in-process)."""
    site_dirs = [d for d in sys.path if d.endswith("site-packages")]
    for site in site_dirs:
        for lib in sorted(glob.glob(os.path.join(site, "nvidia", "*", "lib", "*.so*"))):
            try:
                ctypes.CDLL(lib, mode=ctypes.RTLD_GLOBAL)
            except OSError:
                pass


def fmt_ts(seconds):
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3_600_000)
    m, ms = divmod(ms, 60_000)
    s, ms = divmod(ms, 1000)
    return f"{h:02d}:{m:02d}:{s:02d}.{ms:03d}"


def main():
    if len(sys.argv) != 2:
        print(__doc__.strip(), file=sys.stderr)
        return 2
    media = sys.argv[1]
    if not os.path.isfile(media):
        print(f"Input file not found: {media}", file=sys.stderr)
        return 2

    preload_nvidia_libs()
    try:
        from faster_whisper import WhisperModel
    except ImportError:
        print("faster-whisper is not installed in this Python environment", file=sys.stderr)
        return 3

    model_name = os.environ.get("WHISPER_MODEL", "large-v3-turbo")
    device = os.environ.get("WHISPER_DEVICE")
    attempts = [(device, None)] if device else [("cuda", "float16"), ("cpu", "int8")]

    model = None
    for dev, compute in attempts:
        try:
            model = WhisperModel(model_name, device=dev, compute_type=compute or "default")
            print(f"Whisper model {model_name} loaded on {dev}", flush=True)
            break
        except Exception as err:
            print(f"Could not load model on {dev}: {err}", file=sys.stderr, flush=True)
    if model is None:
        return 1

    segments, info = model.transcribe(media, vad_filter=True)
    print(f"Detected language: {info.language} "
          f"(p={info.language_probability:.2f}), duration {fmt_ts(info.duration)}", flush=True)

    base, _ = os.path.splitext(media)
    vtt_path = f"{base}.{info.language}.vtt"

    lines = ["WEBVTT", ""]
    for seg in segments:
        text = seg.text.strip()
        if not text:
            continue
        lines.append(f"{fmt_ts(seg.start)} --> {fmt_ts(seg.end)}")
        lines.append(text)
        lines.append("")
        # progress for the live log (segments arrive lazily during transcription)
        print(f"[transcribe] {fmt_ts(seg.start)} {text}", flush=True)

    if len(lines) <= 2:
        print("No speech detected — no transcript written", flush=True)
        return 0

    with open(vtt_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print(f"TRANSCRIPT_VTT:{vtt_path}", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())