"""
wivexx_transcribe.py — transcribe AUDIO and VIDEO files to text (+ .srt subtitles).

Reuses Wivexx's already-loaded faster-whisper model (no extra download, no API).
Decoding is handled by PyAV (a faster-whisper dependency that bundles libav/ffmpeg),
so a huge range of containers/codecs works — most audio AND video files: we just
hand the file PATH to model.transcribe() and PyAV demuxes + decodes the audio
stream itself (for a video, it grabs the audio track).

Public:
    is_media(path) -> bool
    media_filetypes() -> [(label, "*.ext ...")]   # for a tk file dialog
    transcribe_file(model, lock, path, language="en",
                    write_txt=True, write_srt=True, on_progress=None) -> dict
"""
import os
import threading

# Broad on purpose — PyAV/libav can demux/decode essentially all of these.
AUDIO_EXTS = {
    ".mp3", ".wav", ".flac", ".m4a", ".m4b", ".aac", ".ogg", ".oga", ".opus",
    ".wma", ".aiff", ".aif", ".aifc", ".alac", ".ape", ".ac3", ".eac3", ".amr",
    ".au", ".caf", ".dts", ".mka", ".mp2", ".mpa", ".ra", ".tta", ".voc", ".wv",
    ".weba", ".3ga", ".gsm", ".spx", ".w64", ".oma", ".dsf", ".dff", ".mlp",
}
VIDEO_EXTS = {
    ".mp4", ".m4v", ".mkv", ".webm", ".mov", ".avi", ".wmv", ".flv", ".f4v",
    ".mpg", ".mpeg", ".mpe", ".m2v", ".mts", ".m2ts", ".ts", ".m2t", ".vob",
    ".3gp", ".3g2", ".ogv", ".rm", ".rmvb", ".asf", ".divx", ".mxf", ".dav",
    ".y4m", ".h264", ".hevc", ".mjpeg", ".qt", ".amv", ".drc", ".gifv", ".ogm",
    ".nut", ".mod", ".tod", ".dvr-ms", ".wtv",
}
MEDIA_EXTS = AUDIO_EXTS | VIDEO_EXTS


def is_media(path):
    return os.path.splitext(path or "")[1].lower() in MEDIA_EXTS


def media_filetypes():
    a = " ".join("*" + e for e in sorted(AUDIO_EXTS))
    v = " ".join("*" + e for e in sorted(VIDEO_EXTS))
    return [("Audio & video", a + " " + v), ("Audio", a), ("Video", v), ("All files", "*.*")]


def _srt_ts(t):
    if t < 0:
        t = 0
    h, rem = divmod(int(t), 3600)
    m, s = divmod(rem, 60)
    ms = int(round((t - int(t)) * 1000))
    if ms >= 1000:
        s += 1
        ms = 0
    return "%02d:%02d:%02d,%03d" % (h, m, s, ms)


def _to_srt(segs):
    out = []
    for i, s in enumerate(segs, 1):
        out.append(str(i))
        out.append("%s --> %s" % (_srt_ts(s.start), _srt_ts(s.end)))
        out.append((s.text or "").strip())
        out.append("")
    return "\n".join(out)


# ── dedicated transcription model (separate from the voice STT) ────────────────
# A general "transcribe any file" job wants the MOST ACCURATE model the machine can
# run AND multilingual auto-detect (Spanish/English/…). The voice loop's model is
# medium.EN (English-only, tuned for speed); using a separate model here also means
# a long transcription never blocks Wivexx from listening.
_model = None
_model_desc = ""
_fast_model = None
_fast_desc = ""
_model_lock = threading.Lock()
_logger = None


def set_logger(logger):
    global _logger
    _logger = logger


def _log(level, msg, *a):
    if _logger:
        getattr(_logger, level, _logger.info)(msg, *a)


# ── GPU guard ─────────────────────────────────────────────────────────────────
# On Windows, ctranslate2 does NOT cleanly raise a CUDA OOM when VRAM is exhausted
# — `WhisperModel(device="cuda")` can HANG inside cudaMalloc, which froze the
# captions/transcribe pipeline forever (no fallback ever ran). So before trying any
# CUDA tier we check the GPU's free VRAM, skip CUDA if there isn't enough headroom,
# and load CUDA under a watchdog thread so an unforeseen stall can't hang the app.
_CREATE_NO_WINDOW = 0x08000000

# Rough free VRAM (MiB) needed to load each model on CUDA in float16, including
# cuDNN workspace + word-timestamp alignment buffers — generous margins on purpose:
# we only touch CUDA when there's COMFORTABLE headroom, because a CUDA allocation
# that runs short doesn't fail cleanly on Windows, it HANGS (see _load_model).
_CUDA_VRAM_NEED = {
    "tiny": 700, "tiny.en": 700, "base": 1000, "base.en": 1000,
    "small": 1800, "small.en": 1800, "medium": 3000, "medium.en": 3000,
    "large-v3-turbo": 3200, "deepdml/faster-whisper-large-v3-turbo-ct2": 3200,
    "mobiuslabsgmbh/faster-whisper-large-v3-turbo": 3200,
    "large": 5000, "large-v1": 5000, "large-v2": 5000, "large-v3": 5000,
}


def _gpu_free_mib():
    """Free VRAM on GPU 0 in MiB, or None if it can't be determined. Tries pynvml,
    then nvidia-smi by name AND by its standard full paths — Wivexx is launched via
    a VBS/pythonw with a minimal PATH where bare `nvidia-smi` may not resolve, so we
    must not depend on PATH (a None reading makes us skip CUDA, costing GPU speed)."""
    try:
        import pynvml
        pynvml.nvmlInit()
        try:
            h = pynvml.nvmlDeviceGetHandleByIndex(0)
            return int(pynvml.nvmlDeviceGetMemoryInfo(h).free // (1024 * 1024))
        finally:
            try:
                pynvml.nvmlShutdown()
            except Exception:
                pass
    except Exception:
        pass
    import subprocess
    root = os.environ.get("SystemRoot", r"C:\Windows")
    candidates = [
        "nvidia-smi",
        os.path.join(root, "System32", "nvidia-smi.exe"),
        r"C:\Program Files\NVIDIA Corporation\NVSMI\nvidia-smi.exe",
    ]
    for exe in candidates:
        try:
            out = subprocess.run(
                [exe, "--query-gpu=memory.free", "--format=csv,noheader,nounits"],
                capture_output=True, timeout=6, creationflags=_CREATE_NO_WINDOW)
            if out.returncode == 0:
                txt = (out.stdout or b"").decode("ascii", "ignore")
                for line in txt.splitlines():
                    line = line.strip()
                    if line:
                        return int(float(line))
        except Exception:
            continue
    return None


def _load_model(name, dev, ct, timeout=None):
    """Construct a WhisperModel. For CUDA, run it under a watchdog thread so a
    stalled cudaMalloc / cuDNN init can't hang the caller — if it doesn't finish
    within `timeout`s we abandon it (the daemon thread is left to die) and raise,
    so get_model() falls through to the CPU tiers."""
    from faster_whisper import WhisperModel
    if dev != "cuda" or not timeout:
        return WhisperModel(name, device=dev, compute_type=ct)
    box = {}

    def _work():
        try:
            box["m"] = WhisperModel(name, device=dev, compute_type=ct)
        except Exception as e:                       # pragma: no cover - device dependent
            box["e"] = e

    th = threading.Thread(target=_work, daemon=True, name="whisper-cuda-load")
    th.start()
    th.join(timeout)
    if th.is_alive():
        raise TimeoutError("CUDA load stalled (>%ss) — VRAM likely exhausted" % timeout)
    if "e" in box:
        raise box["e"]
    return box["m"]


def _load_tiers(tiers):
    """Walk a (name, device, compute_type) preference list and return the first
    (model, "name on dev/ct") that loads. CUDA tiers are skipped when free VRAM is
    too low and loaded under a watchdog, so a saturated GPU degrades to CPU instead
    of hanging. Raises RuntimeError if nothing loads."""
    free = _gpu_free_mib()
    _log("info", "transcribe: GPU free VRAM = %s", "%sMiB" % free if free is not None else "unknown")
    last = None
    seen = set()
    for name, dev, ct in tiers:
        key = (name, dev, ct)
        if key in seen:
            continue
        seen.add(key)
        if dev == "cuda":
            # CONSERVATIVE: only use CUDA with a confident, comfortable free-VRAM
            # reading. Unknown (None) ⇒ skip — CPU is slower but NEVER hangs, whereas
            # a CUDA alloc that runs short hangs the whole feature (it holds the GIL,
            # which can even defeat the load watchdog below).
            need = _CUDA_VRAM_NEED.get(name, 5000)
            if free is None or free < need:
                _log("info", "transcribe: skip %s/cuda — %s free, need ~%sMiB",
                     name, "unknown" if free is None else ("%sMiB" % free), need)
                continue
        try:
            m = _load_model(name, dev, ct, timeout=30 if dev == "cuda" else None)
            return m, "%s on %s/%s" % (name, dev, ct)
        except Exception as e:
            last = e
            _log("info", "transcribe: %s/%s/%s unavailable — %s", name, dev, ct, e)
    raise RuntimeError("no Whisper model could load: %s" % last)


def get_model():
    """Lazy-load the most powerful model this PC can run. Tries large-v3 on the GPU,
    then steps down to whatever loads (so it 'just works' on any machine). CUDA tiers
    are skipped when free VRAM is too low and loaded under a watchdog (see above), so
    a saturated GPU degrades to CPU instead of hanging."""
    global _model, _model_desc
    if _model is not None:
        return _model
    with _model_lock:
        if _model is not None:
            return _model
        try:
            import wivexx_config as _cfg
            preferred = getattr(_cfg, "WHISPER_TRANSCRIBE_MODEL", "large-v3")
        except Exception:
            preferred = "large-v3"
        # most accurate first; multilingual names (no .en) so language auto-detects
        tiers = [
            (preferred, "cuda", "float16"),
            ("large-v3", "cuda", "float16"),
            ("medium", "cuda", "float16"),
            (preferred, "cpu", "int8"),
            ("medium", "cpu", "int8"),
            ("small", "cpu", "int8"),
            ("base", "cpu", "int8"),
        ]
        _model, _model_desc = _load_tiers(tiers)
        _log("info", "transcribe: model ready — %s", _model_desc)
        return _model


def get_fast_model():
    """A lighter, FAST multilingual model for INTERACTIVE jobs (captions editor),
    where waiting on large-v3 hurts. Prefers large-v3-turbo (already cached, ~1.5×
    faster than large-v3 on CPU with near-equal accuracy, full multilingual + word
    timing), then steps down. Same VRAM guard + watchdog as get_model(). Cached
    separately so the standalone transcribe feature keeps its most-accurate model."""
    global _fast_model, _fast_desc
    if _fast_model is not None:
        return _fast_model
    with _model_lock:
        if _fast_model is not None:
            return _fast_model
        tiers = [
            ("large-v3-turbo", "cuda", "float16"),
            ("large-v3-turbo", "cpu", "int8"),
            ("small", "cpu", "int8"),
            ("medium", "cpu", "int8"),
            ("base", "cpu", "int8"),
        ]
        _fast_model, _fast_desc = _load_tiers(tiers)
        _log("info", "captions: fast model ready — %s", _fast_desc)
        return _fast_model


def unload_model():
    """Free the transcription models (e.g. to give VRAM/RAM back before gaming)."""
    global _model, _model_desc, _fast_model, _fast_desc
    with _model_lock:
        if _model is not None or _fast_model is not None:
            _model = None
            _model_desc = ""
            _fast_model = None
            _fast_desc = ""
            try:
                import gc
                gc.collect()
            except Exception:
                pass


def _out_path(src, ext):
    """Write next to the source if we can, else ~/Documents/Wivexx Transcripts."""
    base = os.path.splitext(src)[0]
    d = os.path.dirname(base) or "."
    if os.access(d, os.W_OK):
        return base + ext
    docs = os.path.join(os.path.expanduser("~"), "Documents", "Wivexx Transcripts")
    try:
        os.makedirs(docs, exist_ok=True)
    except Exception:
        pass
    return os.path.join(docs, os.path.basename(base) + ext)


def transcribe_file(path, language=None,
                    write_txt=True, write_srt=True, on_progress=None):
    """Transcribe one media file with the dedicated, most-accurate model.
    language=None → AUTO-DETECT (so Spanish or English just work). Returns a dict
    with text, .txt/.srt paths, detected language and stats."""
    if not os.path.exists(path):
        return {"error": "not found"}
    model = get_model()
    kwargs = dict(beam_size=5, vad_filter=True,
                  vad_parameters={"min_silence_duration_ms": 500})
    if language:
        kwargs["language"] = language             # else Whisper auto-detects the language
    segs = []
    with _model_lock:                             # serialise our own model's use
        segments, info = model.transcribe(path, **kwargs)
        dur = float(getattr(info, "duration", 0.0) or 0.0)
        for s in segments:                        # consuming the generator does the work
            segs.append(s)
            if on_progress and dur > 0:
                try:
                    on_progress(min(1.0, float(s.end) / dur))
                except Exception:
                    pass
    text = " ".join((s.text or "").strip() for s in segs).strip()
    res = {"text": text, "segments": len(segs), "words": len(text.split()),
           "duration": dur, "language": getattr(info, "language", language) or language}
    if write_txt:
        try:
            p = _out_path(path, ".txt")
            with open(p, "w", encoding="utf-8") as f:
                f.write(text + "\n")
            res["txt_path"] = p
        except Exception as e:
            res["txt_error"] = str(e)
    if write_srt and segs:
        try:
            p = _out_path(path, ".srt")
            with open(p, "w", encoding="utf-8") as f:
                f.write(_to_srt(segs))
            res["srt_path"] = p
        except Exception as e:
            res["srt_error"] = str(e)
    return res
