#!/usr/bin/env python3
"""Persistent Piper TTS worker – single subprocess, model loaded once."""

import os
import sys
import json
import wave
import io


class PiperStdoutSink:
    """Wraps stdout so we write exactly N bytes."""

    def __init__(self):
        self._remaining = 0

    def set_length(self, length: int):
        self._remaining = length

    def write(self, data: bytes):
        to_write = min(len(data), self._remaining)
        if to_write > 0:
            sys.stdout.buffer.write(data[:to_write])
            sys.stdout.buffer.flush()
            self._remaining -= to_write


def main():
    from piper import PiperVoice

    model_path = ""
    config_path = None
    voice = None

    # Defaults (will be overridden by init message)
    _default_dir = os.environ.get("PIPER_MODELS_DIR", "/tmp/quibot-piper-models")
    DEFAULT_MODEL = os.path.join(_default_dir, "ca_ES-upc_ona-medium.onnx")
    DEFAULT_CONFIG = os.path.join(_default_dir, "ca_ES-upc_ona-medium.onnx.json")

    # Signal node that the process is alive and listening
    print(json.dumps({"type": "ready"}), flush=True)

    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue

        try:
            msg = json.loads(line)
        except json.JSONDecodeError:
            continue

        if msg.get("type") == "init":
            model_path = msg.get("model", DEFAULT_MODEL) or DEFAULT_MODEL
            config_path = msg.get("config", None)
            print(f"[piper-worker] Loading model='{model_path}'", file=sys.stderr, flush=True)
            try:
                voice = PiperVoice.load(model_path, config_path=config_path)
                print(json.dumps({"type": "init_ok"}), flush=True)
            except Exception as exc:
                err_msg = str(exc).replace('"', '\\"')
                print(json.dumps({"type": "init_error", "error": err_msg}), flush=True)

        elif msg.get("type") == "synthesize":
            text = msg.get("text", "")
            msg_id = msg.get("msgId", "")

            # NEW: output file path from message
            out_path = msg.get("outPath")

            if not voice:
                print(json.dumps({
                    "type": "error",
                    "text": "Model not loaded, send init first",
                    "msgId": msg_id,
                }), flush=True)
                continue

            if not out_path:
                print(json.dumps({
                    "type": "error",
                    "text": "Missing outPath",
                    "msgId": msg_id,
                }), flush=True)
                continue

            try:
                import io
                import wave

                buf = io.BytesIO()
                wf = wave.open(buf, 'wb')

                voice.synthesize_wav(text, wf)
                wf.close()

                wav_bytes = buf.getvalue()

                # Write to file instead of stdout
                os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
                with open(out_path, "wb") as f:
                    f.write(wav_bytes)

                print(json.dumps({
                    "type": "synthesized",
                    "bytes": len(wav_bytes),
                    "msgId": msg_id,
                    "outPath": out_path
                }), flush=True)

            except Exception as exc:
                err_msg = str(exc).replace('"', '\\"')
                print(json.dumps({
                    "type": "error",
                    "text": err_msg,
                    "msgId": msg_id,
                }), flush=True)

if __name__ == "__main__":
    main()
