#!/usr/bin/env python3
"""MeloTTS server for inferon.

Long-running subprocess: loads the MeloTTS model once, then reads lines of
text on stdin and writes "<wav-path> <duration-seconds>" per line on stdout.
stderr passes through for logging.

Protocol:
    -> one line of text to synthesize
    <- e.g. "/tmp/inferon/tts-1723.wav 3.42"

Kill with SIGTERM — no cleanup needed, WAVs live in /tmp.
"""

import sys
import os
import time
import signal
import warnings

warnings.filterwarnings("ignore")

MELO_ROOT = os.path.expanduser("~/Projects/MeloTTS")
CKPT = os.path.join(MELO_ROOT, "ckpt/MeloTTS-English-v3/checkpoint.pth")
CONFIG = os.path.join(MELO_ROOT, "ckpt/MeloTTS-English-v3/config.json")
OUT_DIR = "/tmp/inferon"
SPEAKER = "EN-Newest"  # v3 English: single speaker (British-accented female)
LANG = "EN"
SPEED = 1.3  # speech rate; 1.0 = natural, higher = faster

sys.path.insert(0, MELO_ROOT)


def main() -> None:
    signal.signal(signal.SIGTERM, lambda *_: sys.exit(0))
    signal.signal(signal.SIGINT, lambda *_: sys.exit(0))

    os.makedirs(OUT_DIR, exist_ok=True)

    # Heavy imports inside main so signal handlers are installed first.
    from melo.api import TTS

    model = TTS(language=LANG, config_path=CONFIG, ckpt_path=CKPT)

    speaker_ids = model.hps.data.spk2id
    if SPEAKER not in speaker_ids:
        # English v3 fallback: use the first speaker and say so on stderr.
        fallback = next(iter(speaker_ids))
        print(f"melo_server: speaker {SPEAKER!r} not in {list(speaker_ids)}; using {fallback!r}",
              file=sys.stderr, flush=True)
        speaker = fallback
    else:
        speaker = speaker_ids[SPEAKER]

    # Warm up with a full sentence (longer text paths hit NLTK/num2words,
    # which is where missing-data errors surface). Fail before READY.
    warmup = "/tmp/inferon/tts-warmup.wav"
    model.tts_to_file("Warmup sentence number twelve, spoken on the first of January.",
                      speaker, warmup, speed=SPEED)
    os.unlink(warmup)

    # Ready marker — inferon waits for this before declaring the backend up.
    print("READY", flush=True)

    for line in sys.stdin:
        text = line.strip()
        if not text:
            continue
        text = text[:2000]  # protocol cap — extremely long replies get clipped

        wav_path = os.path.join(OUT_DIR, f"tts-{time.time_ns()}.wav")
        try:
            start = time.monotonic()
            model.tts_to_file(text, speaker, wav_path, speed=SPEED)
            dur = time.monotonic() - start
            print(f"{wav_path} {dur:.2f}", flush=True)
        except Exception as e:  # noqa: BLE001 - protocol: report and survive
            print(f"melo_server: synthesis failed: {e}", file=sys.stderr, flush=True)
            print("ERROR", flush=True)


if __name__ == "__main__":
    main()