"""generate_speech over MCP streamable HTTP (media-node, the author's home GPU node with qwentts.cpp).

Needs a server with the ru-stress model and a saved project voice (tts.media_node.voice, e.g. made with
its save_voice tool from the reference clip). Batches of up to 64 texts / 29000 characters per call.
"""
import itertools
import json
import urllib.request

_ids = itertools.count(1)


class MCP:
    def __init__(self, url):
        self.url, self.sid = url, None
        self._post({"jsonrpc": "2.0", "id": next(_ids), "method": "initialize",
                    "params": {"protocolVersion": "2025-06-18", "capabilities": {},
                               "clientInfo": {"name": "ensvoice", "version": "0.4"}}})
        self._post({"jsonrpc": "2.0", "method": "notifications/initialized"}, expect=False)

    def _post(self, msg, expect=True, timeout=7200):
        h = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
        if self.sid:
            h["Mcp-Session-Id"] = self.sid
        with urllib.request.urlopen(urllib.request.Request(self.url, json.dumps(msg).encode(), h), timeout=timeout) as r:
            self.sid = r.headers.get("Mcp-Session-Id", self.sid)
            body = r.read().decode()
        if not expect:
            return None
        if body.lstrip().startswith("{"):
            return json.loads(body)
        last = None
        for line in body.splitlines():
            if line.startswith("data:"):
                d = json.loads(line[5:])
                if d.get("id") == msg.get("id"):
                    last = d
        return last

    def call(self, name, args):
        r = self._post({"jsonrpc": "2.0", "id": next(_ids), "method": "tools/call",
                        "params": {"name": name, "arguments": args}})
        if "error" in r:
            raise RuntimeError(r["error"])
        return r["result"]


class MediaNodeTTS:
    def __init__(self, cfg):
        c = cfg["tts"]["media_node"]
        self.url, self.voice = c["url"], c["voice"]
        self.rewrite = c.get("download_host_rewrite")  # ["from", "to"] if file URLs name another host
        self.m = MCP(self.url)

    def speak(self, items, seed, log=print):
        batches, cur, c = [], [], 0
        for it in items:
            if cur and (len(cur) == 64 or c + len(it[0]) > 29000):
                batches.append(cur); cur, c = [], 0
            cur.append(it); c += len(it[0])
        if cur:
            batches.append(cur)
        for bi, b in enumerate(batches):
            for attempt in range(3):
                try:
                    r = self.m.call("generate_speech", {"model": "ru-stress", "voice": self.voice, "language": "Russian",
                                                        "seed": seed, "texts": [{"text": t} for t, _ in b]})
                    break
                except Exception as ex:  # a backend timeout: once more with a new session
                    log(f"retry {bi}: {ex}"); self.m = MCP(self.url)
            files = sorted(r["structuredContent"]["files"], key=lambda f: f["item"])
            for (_, out), f in zip(b, files):
                url = f["url"] if not self.rewrite else f["url"].replace(*self.rewrite)
                urllib.request.urlretrieve(url, out)
            log(f"seed {seed}: batch {bi + 1}/{len(batches)}")