"""Persistent MeloTTS worker (Korean).

MeloTTS lives in its own Python (melo311); loading the model takes seconds, so
we load it ONCE here and then serve synthesis requests over stdin/stdout. This
process is launched with the melo311 interpreter by wsai.backends.melo.MeloTTS.

MeloTTS (and its deps) print progress straight to stdout, which would corrupt
the JSON protocol. So on startup we split the streams: a private duplicate of
the original stdout carries the protocol, and fd 1 is redirected to fd 2 so all
library chatter lands on stderr instead.

Protocol (one JSON object per line, on the protocol channel):
    <- {"text": "...", "out": "/abs/path.wav", "speed": 1.3}
    -> {"ok": true, "out": "/abs/path.wav", "ms": 123}
    -> {"ok": false, "error": "..."}
On startup, once the model is ready, it emits exactly one line:
    -> {"ready": true, "ms": <load-ms>, "device": "cpu"}
"""

import json
import os
import sys
import time

# Split protocol from library noise BEFORE importing anything heavy.
_proto = os.fdopen(os.dup(1), "w", buffering=1)  # private copy of real stdout
os.dup2(2, 1)  # fd1 -> stderr, so stray library prints don't hit the protocol


def _emit(obj: dict) -> None:
    _proto.write(json.dumps(obj) + "\n")
    _proto.flush()


def _log(*a):
    print(*a, file=sys.stderr, flush=True)


def main() -> None:
    lang = "KR"
    requested = os.environ.get("WSAI_MELO_DEVICE", "auto")  # cpu | cuda | auto
    from melo.api import TTS  # heavy import; only in the melo venv

    def _has_cuda() -> bool:
        try:
            import torch

            return torch.cuda.is_available()
        except Exception:
            return False

    device = requested
    if requested == "auto":
        device = "cuda" if _has_cuda() else "cpu"

    t0 = time.monotonic()
    try:
        tts = TTS(language=lang, device=device)
    except Exception as exc:
        # CUDA picked but unusable (CPU-only torch, missing libs, OOM): fall back
        # to CPU rather than leaving the whole voice loop dead.
        if device == "cuda":
            _log(f"[melo_worker] CUDA load failed ({exc}); falling back to CPU")
            device = "cpu"
            tts = TTS(language=lang, device=device)
        else:
            raise
    speaker_id = tts.hps.data.spk2id[lang]
    load_ms = int((time.monotonic() - t0) * 1000)
    _emit({"ready": True, "ms": load_ms, "device": device})
    _log(f"[melo_worker] model ready in {load_ms} ms on {device}")

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            req = json.loads(line)
            text = req["text"]
            out = req["out"]
            speed = float(req.get("speed", 1.0))
            if out.startswith("/tmp") or out.startswith("/dev/shm"):
                raise ValueError(f"refusing RAM-backed tmpfs path: {out}")
            s = time.monotonic()
            tts.tts_to_file(text, speaker_id, out, speed=speed)
            ms = int((time.monotonic() - s) * 1000)
            _emit({"ok": True, "out": out, "ms": ms})
        except Exception as exc:  # keep the worker alive across bad requests
            _emit({"ok": False, "error": f"{type(exc).__name__}: {exc}"})
            _log(f"[melo_worker] error: {exc}")


if __name__ == "__main__":
    main()