#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
VDQ transcribe worker — un seul script, trois modes, déployé sur boxi ET AIO.

1) Mode job-file (boxi fallback local) :
     VDQ_RUNNER=boxi python worker.py <job.json>
   job = {"wav","out_dir","base","model","result"} (chemins locaux) →
   écrit <base>.srt/.txt/.transcript.json dans out_dir + fichier result.

2) Mode --serve (AIO, transport par ssh, PAS de SMB) :
     cat video.wav | ssh aio "python worker.py --serve <id> <model>"
   lit le wav sur stdin, transcrit, écrit dans un dossier temp local AIO,
   puis imprime en DERNIÈRE ligne : @@VDQ_B64@@<base64 json>
   où json = {"status","meta...","files":{"srt","txt","tjson"}} (contenus b64).

3) Mode --write-files (boxi, finalisation) :
     printf '%s' "<ligne @@VDQ_B64@@>" | python worker.py --write-files <outdir> <stem>
   décode l'enveloppe et écrit <stem>.srt/.txt/.transcript.json dans <outdir>.

Config : VAD off (garbage sur la musique, bench 2026-10-16),
condition_on_previous_text=False (anti-boucle docs longs), beam_size=1 (vitesse CPU).
"""
import base64
import json
import os
import sys
import time


def norm_path(p):
    """Windows : //server/share en forward slashes n'est pas un UNC valide → backslashes."""
    if os.name == "nt" and p.startswith("//"):
        return p.replace("/", "\\")
    return p


def setup_cuda_dlls():
    """Windows : rendre les DLL CUDA (cuBLAS/cuDNN/cudart) trouvables par ctranslate2.
    Sources : wheels pip nvidia-*-cu12 (site-packages/nvidia/*/bin)."""
    if os.name != "nt":
        return
    import sysconfig
    sp = sysconfig.get_paths()["purelib"]
    for d in ("cublas", "cudnn", "cuda_runtime"):
        p = os.path.join(sp, "nvidia", d, "bin")
        if os.path.isdir(p):
            try:
                os.add_dll_directory(p)
                os.environ["PATH"] = p + os.pathsep + os.environ.get("PATH", "")
            except Exception:
                pass


def pick_device():
    """cuda si un GPU est visible (override VDQ_FORCE_CPU=1), sinon cpu.
    Retourne (device, compute_type)."""
    if os.environ.get("VDQ_FORCE_CPU") == "1":
        return "cpu", "int8"
    try:
        import ctranslate2
        if ctranslate2.get_cuda_device_count() > 0:
            return "cuda", "float16"
    except Exception:
        pass
    return "cpu", "int8"


setup_cuda_dlls()


def fmt_ts(seconds):
    ms = int(round(seconds * 1000))
    h, ms = divmod(ms, 3600000)
    m, ms = divmod(ms, 60000)
    s, ms = divmod(ms, 1000)
    return "%02d:%02d:%02d,%03d" % (h, m, s, ms)


def write_srt(path, segments):
    lines = []
    for i, seg in enumerate(segments, 1):
        lines.append("%d" % i)
        lines.append("%s --> %s" % (fmt_ts(seg["start"]), fmt_ts(seg["end"])))
        lines.append(seg["text"].strip())
        lines.append("")
    with open(path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))


def transcribe_to_files(wav, out_dir, base, model_name, device=None, ctype=None):
    """Retourne (segments, meta_dict). Écrit srt/txt/json dans out_dir.
    GPU (cuda float16) tenté d'abord — bench 2026-10-16 : 10,9× realtime sur RTX 3070
    vs 3× (i7 CPU) vs 1× (boxi CPU). Fallback CPU int8 si init/VRAM échoue."""
    from faster_whisper import WhisperModel
    if device is None:
        device, ctype = pick_device()
    print("[worker] loading model %s (%s, %s)..." % (model_name, device, ctype), flush=True)
    try:
        model = WhisperModel(model_name, device=device, compute_type=ctype)
    except Exception as e:
        if device == "cuda":
            print("[worker] CUDA init échoué (%s) — fallback CPU" % str(e)[:200], flush=True)
            device, ctype = "cpu", "int8"
            model = WhisperModel(model_name, device="cpu", compute_type="int8")
        else:
            raise
    print("[worker] transcribing %s..." % wav, flush=True)
    segments_iter, info = model.transcribe(
        wav,
        language=None,
        vad_filter=False,
        beam_size=1,
        condition_on_previous_text=False,
    )
    segments = []
    for seg in segments_iter:
        segments.append({
            "start": round(float(seg.start), 3),
            "end": round(float(seg.end), 3),
            "text": seg.text.strip(),
        })
    srt_path = os.path.join(out_dir, base + ".srt")
    txt_path = os.path.join(out_dir, base + ".txt")
    json_path = os.path.join(out_dir, base + ".transcript.json")
    full_text = " ".join(s["text"] for s in segments).strip()
    write_srt(srt_path, segments)
    with open(txt_path, "w", encoding="utf-8") as f:
        f.write(full_text + "\n")
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump({
            "language": info.language,
            "language_probability": round(float(info.language_probability), 3),
            "duration": round(float(info.duration), 2),
            "device": device,
            "segments": segments,
        }, f, ensure_ascii=False)
    meta = {
        "language": info.language,
        "segments": len(segments),
        "audio_seconds": round(float(info.duration), 1),
        "device": device,
    }
    print("[worker] done: %d segments, %.1fs audio (device=%s)" % (len(segments), float(info.duration), device), flush=True)
    return segments, meta


def read_stdin_to(path, chunk=1 << 20):
    with open(path, "wb") as f:
        while True:
            data = sys.stdin.buffer.read(chunk)
            if not data:
                break
            f.write(data)
    return os.path.getsize(path)


def cleanup_old_temp(tempdir, max_age_s=86400):
    now = time.time()
    try:
        for name in os.listdir(tempdir):
            p = os.path.join(tempdir, name)
            try:
                if now - os.path.getmtime(p) > max_age_s:
                    if os.path.isdir(p):
                        import shutil
                        shutil.rmtree(p, ignore_errors=True)
                    else:
                        os.remove(p)
            except OSError:
                pass
    except OSError:
        pass


def mode_serve(job_id, model_name):
    tempdir = os.path.join(os.path.expanduser("~"), "DevTemp", "vdq-transcribe")
    os.makedirs(tempdir, exist_ok=True)
    cleanup_old_temp(tempdir)
    wav = os.path.join(tempdir, job_id + ".wav")
    out_dir = os.path.join(tempdir, job_id + ".out")
    os.makedirs(out_dir, exist_ok=True)
    size = read_stdin_to(wav)
    print("[worker] wav received: %d bytes" % size, flush=True)
    try:
        t0 = time.time()
        segments, meta = transcribe_to_files(wav, out_dir, job_id, model_name)
        meta["elapsed_s"] = round(time.time() - t0, 1)
        files = {}
        for key, ext in (("srt", ".srt"), ("txt", ".txt"), ("tjson", ".transcript.json")):
            with open(os.path.join(out_dir, job_id + ext), "rb") as f:
                files[key] = base64.b64encode(f.read()).decode("ascii")
        envelope = {"status": "done", "meta": meta, "files": files}
        print("@@VDQ_B64@@" + base64.b64encode(
            json.dumps(envelope, ensure_ascii=False).encode("utf-8")).decode("ascii"), flush=True)
    except Exception as e:
        print("@@VDQ_B64@@" + base64.b64encode(json.dumps(
            {"status": "error", "error": str(e)[:500]}).encode("utf-8")).decode("ascii"), flush=True)
        sys.exit(1)
    finally:
        try:
            os.remove(wav)
        except OSError:
            pass


def mode_write_files(out_dir, stem):
    data = sys.stdin.read()
    marker = "@@VDQ_B64@@"
    idx = data.rfind(marker)
    if idx < 0:
        print("[worker] ERROR: pas d'enveloppe @@VDQ_B64@@ sur stdin")
        sys.exit(1)
    envelope = json.loads(base64.b64decode(data[idx + len(marker):]).decode("utf-8"))
    if envelope.get("status") != "done":
        print("[worker] ERROR: %s" % envelope.get("error", "enveloppe non-done"))
        sys.exit(1)
    files = envelope["files"]
    mapping = (("srt", ".srt"), ("txt", ".txt"), ("tjson", ".transcript.json"))
    for key, ext in mapping:
        with open(os.path.join(out_dir, stem + ext), "wb") as f:
            f.write(base64.b64decode(files[key]))
    meta = envelope.get("meta", {})
    print("[worker] files written: %s{.srt,.txt,.transcript.json} lang=%s segs=%s elapsed=%ss device=%s" % (
        stem, meta.get("language"), meta.get("segments"), meta.get("elapsed_s"), meta.get("device")), flush=True)


def mode_job_file(job_path):
    runner = os.environ.get("VDQ_RUNNER", "boxi")
    with open(norm_path(job_path), "r", encoding="utf-8") as f:
        job = json.load(f)
    result_path = norm_path(job["result"])
    base = job["base"]
    out_dir = norm_path(job["out_dir"]).rstrip("/\\")
    wav = norm_path(job["wav"])
    model_name = job.get("model", "small")

    def write_result(payload):
        payload["runner"] = runner
        payload["model"] = model_name
        try:
            with open(result_path, "w", encoding="utf-8") as f:
                json.dump(payload, f, ensure_ascii=False)
        except Exception as e:
            print("[worker] cannot write result: %s" % e)

    if not os.path.isfile(wav):
        write_result({"status": "error", "error": "wav introuvable: %s" % wav})
        sys.exit(1)
    try:
        t0 = time.time()
        segments, meta = transcribe_to_files(wav, out_dir, base, model_name)
        meta["elapsed_s"] = round(time.time() - t0, 1)
        write_result({"status": "done", **meta})
    except Exception as e:
        write_result({"status": "error", "error": str(e)[:500]})
        print("[worker] ERROR: %s" % e, flush=True)
        sys.exit(1)


def main():
    args = sys.argv[1:]
    if args and args[0] == "--serve" and len(args) >= 3:
        mode_serve(args[1], args[2])
    elif args and args[0] == "--write-files" and len(args) >= 3:
        mode_write_files(args[1], args[2])
    elif args:
        mode_job_file(args[0])
    else:
        print("usage: worker.py <job.json> | --serve <id> <model> | --write-files <outdir> <stem>")
        sys.exit(2)


if __name__ == "__main__":
    main()