"""zone transcribe worker - runs inside the venv created by `zone transcribe setup`.

Invocation: python worker.py '<json>'
Protocol: one JSON event per stdout line (progress | result | error).
Exit codes: 0 ok, 1 error. Structured errors only - never a bare traceback.
"""

import json
import os
import sys

if hasattr(sys.stdout, "reconfigure"):
    sys.stdout.reconfigure(encoding="utf-8")


def emit(obj):
    sys.stdout.write(json.dumps(obj, ensure_ascii=False) + "\n")
    sys.stdout.flush()


def fail(message, code="WORKER_ERROR", hint=None):
    emit({"event": "error", "code": code, "message": message, "hint": hint})
    sys.exit(1)


def prepare_dlls():
    """Windows: CTranslate2 needs cuDNN/cuBLAS DLLs; the torch wheels bundle them
    in torch/lib. Must run BEFORE importing faster_whisper/ctranslate2."""
    try:
        import torch  # also loads its own DLLs into the process
    except ImportError:
        fail("torch is not installed in the worker venv.", "NOT_READY",
             'Run "zone transcribe setup".')
    lib = os.path.join(os.path.dirname(torch.__file__), "lib")
    if os.name == "nt" and os.path.isdir(lib):
        os.add_dll_directory(lib)
        os.environ["PATH"] = lib + os.pathsep + os.environ.get("PATH", "")
    return torch


def import_faster_whisper():
    try:
        import faster_whisper
        return faster_whisper
    except ImportError:
        fail("faster-whisper is not installed in the worker venv.", "NOT_READY",
             'Run "zone transcribe setup".')


def cached_models():
    hub = os.path.join(os.environ.get("HF_HOME", ""), "hub")
    if not os.path.isdir(hub):
        return []
    names = []
    for d in os.listdir(hub):
        if d.startswith("models--") and "whisper" in d.lower():
            names.append(d.replace("models--", "").replace("--", "/"))
    return sorted(names)


def probe():
    torch = prepare_dlls()
    faster_whisper = import_faster_whisper()
    try:
        import ctranslate2
    except ImportError:
        fail("ctranslate2 is not installed in the worker venv.", "NOT_READY",
             'Run "zone transcribe setup".')
    cuda = bool(torch.cuda.is_available())
    emit({
        "event": "result",
        "python": sys.version.split()[0],
        "torch": torch.__version__,
        "cuda_available": cuda,
        "cuda_device": torch.cuda.get_device_name(0) if cuda else None,
        "vram_mb": round(torch.cuda.get_device_properties(0).total_memory / 1048576) if cuda else None,
        "faster_whisper": faster_whisper.__version__,
        "ctranslate2": ctranslate2.__version__,
        "models_cached": cached_models(),
    })


def pull(model):
    prepare_dlls()
    download_model = import_faster_whisper().download_model
    emit({"event": "progress", "stage": "download", "message": f"Downloading {model}..."})
    download_model(model)
    emit({"event": "result", "model": model, "downloaded": True})


LADDER_CUDA = [("large-v3", "int8_float16"), ("medium", "int8_float16")]
LADDER_CPU = [("large-v3", "int8")]

DIARIZATION_CHECKPOINTS = [
    "pyannote/speaker-diarization-community-1",
    "pyannote/speaker-diarization-3.1",
]
DIARIZATION_TERMS_HINT = (
    "Accept the model terms at https://huggingface.co/pyannote/speaker-diarization-community-1 "
    "or https://huggingface.co/pyannote/speaker-diarization-3.1 (matching your token's access)."
)


def load_diarization_pipeline():
    if not os.environ.get("HF_TOKEN"):
        fail("HuggingFace token not set.", "NOT_READY",
             'Run "zone transcribe setup" to store one. ' + DIARIZATION_TERMS_HINT)
    try:
        from pyannote.audio import Pipeline
    except ImportError:
        fail("pyannote.audio is not installed in the worker venv.", "NOT_READY",
             'Run "zone transcribe setup".')
    last_err = None
    for checkpoint in DIARIZATION_CHECKPOINTS:
        try:
            emit({"event": "progress", "stage": "diarize",
                  "message": f"Loading diarization pipeline ({checkpoint})..."})
            return Pipeline.from_pretrained(checkpoint)
        except Exception as e:
            last_err = e
    fail(f"Could not load a diarization pipeline: {type(last_err).__name__}: {str(last_err)[:300]}",
         "WORKER_ERROR", DIARIZATION_TERMS_HINT)


def diarize_segments(pipeline, input_path, device, segments):
    """Second pass: pyannote speaker turns merged onto Whisper segments by max overlap.
    Returns the number of distinct speakers labeled."""
    import torch
    if device == "cuda":
        try:
            pipeline.to(torch.device("cuda"))
        except Exception:
            emit({"event": "progress", "stage": "diarize",
                  "message": "CUDA unavailable for diarization - using CPU."})
    from faster_whisper.audio import decode_audio
    audio = decode_audio(input_path, sampling_rate=16000)
    waveform = torch.from_numpy(audio).unsqueeze(0)
    emit({"event": "progress", "stage": "diarize", "message": "Diarizing..."})
    diarization = pipeline({"waveform": waveform, "sample_rate": 16000})
    # pyannote.audio 4.x wraps the Annotation in an output dataclass; 3.x returns it directly.
    annotation = getattr(diarization, "speaker_diarization", diarization)
    turns = [(turn.start, turn.end, label)
             for turn, _, label in annotation.itertracks(yield_label=True)]
    names = {}
    for seg in segments:
        best, best_overlap = None, 0.0
        for (ts, te, label) in turns:
            overlap = min(seg["end"], te) - max(seg["start"], ts)
            if overlap > best_overlap:
                best, best_overlap = label, overlap
        if best is not None:
            if best not in names:
                names[best] = f"Speaker {len(names) + 1}"
            seg["speaker"] = names[best]
    return len(names)


def load_model(name, device, compute_type):
    WhisperModel = import_faster_whisper().WhisperModel
    emit({"event": "progress", "stage": "load",
          "message": f"Loading {name} ({compute_type}) on {device}..."})
    return WhisperModel(name, device=device, compute_type=compute_type)


def transcribe(args):
    torch = prepare_dlls()

    device = args.get("device", "auto")
    forced_cuda = device == "cuda"
    if device == "auto":
        device = "cuda" if torch.cuda.is_available() else "cpu"
    elif forced_cuda and not torch.cuda.is_available():
        fail("CUDA requested but not available in this venv.", "NOT_READY",
             'Re-run "zone transcribe setup" on a machine with an NVIDIA GPU, or use --device cpu.')

    do_diarize = bool(args.get("diarize"))
    dia_pipeline = load_diarization_pipeline() if do_diarize else None

    requested = args.get("model", "auto")
    if requested != "auto":
        ladder = [(requested, "int8_float16" if device == "cuda" else "int8")]
    elif device == "cuda":
        ladder = LADDER_CUDA
    else:
        ladder = LADDER_CPU

    model, used, last_err = None, None, None
    for name, compute_type in ladder:
        try:
            model = load_model(name, device, compute_type)
            used = (name, compute_type)
            break
        except Exception as e:  # OOM or CUDA init failure - walk the ladder
            last_err = e
            if device == "cuda":
                emit({"event": "progress", "stage": "fallback",
                      "message": f"{name} failed on CUDA ({str(e)[:200]}); trying next option..."})
                continue
            raise

    if model is None and device == "cuda":
        if forced_cuda:
            fail(f"CUDA transcription failed: {last_err}", "WORKER_ERROR",
                 "Try --device cpu or a smaller --model.")
        emit({"event": "progress", "stage": "fallback",
              "message": "CUDA unusable - falling back to CPU (slower)."})
        device = "cpu"
        name, compute_type = LADDER_CPU[0] if requested == "auto" else (requested, "int8")
        model = load_model(name, device, compute_type)
        used = (name, compute_type)
    if model is None:
        fail(f"Could not load a Whisper model: {last_err}")

    lang = args.get("lang", "auto")
    emit({"event": "progress", "stage": "transcribe", "message": "Transcribing..."})
    segments_iter, info = model.transcribe(
        args["input"],
        language=None if lang == "auto" else lang,
        vad_filter=True,
    )

    segments = []
    next_report = 60.0
    for seg in segments_iter:
        segments.append({"start": round(seg.start, 3), "end": round(seg.end, 3), "text": seg.text})
        if info.duration and seg.end >= next_report:
            pct = min(99, round(seg.end / info.duration * 100))
            emit({"event": "progress", "stage": "transcribe",
                  "message": f"{pct}% ({round(seg.end)}s / {round(info.duration)}s)"})
            next_report += 60.0

    speakers = None
    if do_diarize:
        # 4 GB VRAM cannot hold Whisper large-v3 and pyannote together - free first.
        del model
        import gc
        gc.collect()
        if device == "cuda":
            torch.cuda.empty_cache()
        speakers = diarize_segments(dia_pipeline, args["input"], device, segments)

    emit({
        "event": "result",
        "segments": segments,
        "meta": {
            "model_requested": requested,
            "model_used": used[0],
            "compute_type": used[1],
            "device": device,
            "downgraded": requested == "auto" and used[0] != "large-v3",
            "duration": round(info.duration, 3) if info.duration is not None else None,
            "language": info.language,
            "language_probability": round(info.language_probability, 3) if info.language_probability is not None else None,
            "diarized": do_diarize,
            "speakers": speakers,
        },
    })


def main():
    if len(sys.argv) < 2:
        fail("Missing JSON argument.", "BAD_INPUT")
    try:
        args = json.loads(sys.argv[1])
    except json.JSONDecodeError as e:
        fail(f"Bad JSON argument: {e}", "BAD_INPUT")

    if not isinstance(args, dict):
        fail("JSON argument must be an object.", "BAD_INPUT")

    mode = args.get("mode")
    try:
        if mode == "probe":
            probe()
        elif mode == "pull":
            pull(args.get("model", "large-v3"))
        elif mode == "transcribe":
            if not args.get("input") or not os.path.isfile(args["input"]):
                fail(f"Input not found: {args.get('input')}", "BAD_INPUT")
            transcribe(args)
        else:
            fail(f"Unknown mode: {mode}", "BAD_INPUT")
    except SystemExit:
        raise
    except Exception as e:  # any unexpected library failure -> structured error
        fail(f"{type(e).__name__}: {e}")


if __name__ == "__main__":
    main()