"""Persistent local GPU voice worker. JSON lines on stdin/stdout; logs on stderr."""
import contextlib
import json
from pathlib import Path
import sys


def main():
    protocol = sys.stdout
    with contextlib.redirect_stdout(sys.stderr):
        import numpy as np
        import soundfile as sf
        import torch
        from chatterbox.tts_turbo import ChatterboxTurboTTS
        torch.set_num_threads(4)
        if not torch.cuda.is_available():
            raise RuntimeError("Chatterbox lab requires the configured NVIDIA GPU")
        model = ChatterboxTurboTTS.from_pretrained(device="cuda")
        model.prepare_conditionals(sys.argv[1])
    print(json.dumps({"ready": True}), file=protocol, flush=True)
    for line in sys.stdin:
        try:
            request = json.loads(line)
            with contextlib.redirect_stdout(sys.stderr):
                # Sentence-sized chunks avoid Turbo's long-utterance truncation.
                import re
                chunks = re.split(r"(?<=[.!?])\s+", request["text"].strip())
                parts = []
                for chunk in chunks:
                    if not chunk:
                        continue
                    # Bound even punctuation-free custom hints.
                    words = chunk.split()
                    for offset in range(0, len(words), 45):
                        samples = model.generate(" ".join(words[offset:offset+45])).squeeze().detach().cpu().numpy()
                        if parts:
                            parts.append(np.zeros(int(model.sr * .18), dtype=np.float32))
                        parts.append(samples)
                if not parts:
                    raise ValueError("Hint text is empty")
                output = Path(request["output"])
                temp = output.with_suffix(".partial")
                sf.write(temp, np.concatenate(parts), model.sr, format="WAV", subtype="PCM_16")
                temp.replace(output)
            print(json.dumps({"ok": True}), file=protocol, flush=True)
        except Exception as error:
            print(json.dumps({"ok": False, "error": str(error)}), file=protocol, flush=True)


if __name__ == "__main__":
    main()