"""Real-time streaming TTS server — FastAPI + WebSocket + Web Audio.

Client sends {text}; server runs the frontend (text->phone/tone/lang ids), then the
streaming vocoder, pushing 16-bit PCM chunks over the WebSocket as each 24-frame
chunk (~384 ms audio) is produced. The browser plays them via Web Audio as they
arrive → first sound in ~tens of ms, not after the whole utterance.

Backend: the sherpa-onnx fork's OfflineTtsMbistftStreamModel (C++/ORT) when available;
falls back to the pure-onnxruntime split runner (bit-exact) otherwise.
"""
import os, json, struct, asyncio
from pathlib import Path

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
import numpy as np

HERE = Path(__file__).resolve().parent
os.environ.setdefault("NLTK_DATA", str(HERE / "nltk_data"))
# v2.1 streaming split (3-voice) is hosted in the model repo to avoid the Space's 1 GB LFS
# limit; downloaded (and cached) at startup. Local override via ENC_ONNX/DEC_ONNX.
def _model(fn):
    from huggingface_hub import hf_hub_download
    return hf_hub_download(os.environ.get("PRIMETTS_REPO", "Luigi/PrimeTTS"), f"v21_streaming/{fn}")
ENC = os.environ.get("ENC_ONNX") or _model("v21_enc.onnx")
DEC = os.environ.get("DEC_ONNX") or _model("v21_dec.onnx")
SR = 16000
CHUNK, LEFT, RIGHT, HOP, CHAN = 24, 64, 16, 256, 192  # non-causal clean vocoder

# ---- frontend (text -> phone/tone/lang ids, add_blank) ----
import sys
sys.path.insert(0, str(HERE))
import frontend_bopomofo as F  # bopomofo + arpabet, 88 syms / 6 tones / 2 langs
F.text_to_ids("您好")  # warm the g2pw model


def _blank(seq):
    o = [0] * (2 * len(seq) + 1)
    o[1::2] = seq
    return o


import re as _re
# Split AFTER punctuation (keep the delimiter with its clause), then greedily merge
# tiny fragments so each chunk is a reasonable phrase (fast enc, natural prosody).
_CLAUSE_RE = _re.compile(r'(?<=[。！？；;!?\n,，、:：])')


def _split_clauses(text, min_chars=12):
    parts = [p for p in _CLAUSE_RE.split(text) if p.strip()]
    out, cur = [], ""
    for p in parts:
        cur += p
        if len(cur.strip()) >= min_chars:
            out.append(cur); cur = ""
    if cur.strip():
        out.append(cur)
    return out or [text]


# ---- backend: sherpa-onnx C++ (preferred) or onnxruntime split runner ----
class Backend:
    def __init__(self):
        self.kind = None
        try:
            try:
                from sherpa_onnx import OfflineTtsMbistftStreamModel as M
            except Exception:
                from _sherpa_onnx import OfflineTtsMbistftStreamModel as M
            self.m = M(enc=ENC, dec=DEC, num_threads=2, right_lookahead=16)
            self.kind = "sherpa"
        except Exception as e:
            print(f"[backend] sherpa-onnx unavailable ({e}); using onnxruntime split runner")
            import onnxruntime as ort
            so = ort.SessionOptions(); so.intra_op_num_threads = 2; so.inter_op_num_threads = 1
            self.enc = ort.InferenceSession(ENC, so, providers=["CPUExecutionProvider"])
            self.dec = ort.InferenceSession(DEC, so, providers=["CPUExecutionProvider"])
            self.kind = "onnx"
        print(f"[backend] {self.kind}")

    def stream(self, phone, tone, lang, noise_scale, length_scale, emit, sid=0):
        """Call emit(pcm_bytes) per audio chunk. emit returns False to stop. sid selects voice."""
        x, tn, lg = _blank(phone), _blank(tone), _blank(lang)
        if self.kind == "sherpa":
            def cb(samples, progress):
                return emit(_pcm16(np.asarray(samples, np.float32)))
            self.m.generate(x=x, tone=tn, lang=lg, noise_scale=noise_scale,
                            length_scale=length_scale, callback=cb)
            return
        # onnxruntime split: enc once (with speaker sid) + chunked overlap-save dec
        z = self.enc.run(None, {
            "x": np.array([x], np.int64), "tone": np.array([tn], np.int64),
            "lang": np.array([lg], np.int64), "x_lengths": np.array([len(x)], np.int64),
            "noise_scale": np.array([noise_scale], np.float32),
            "length_scale": np.array([length_scale], np.float32),
            "sid": np.array([sid], np.int64)})[0]
        Tf = z.shape[2]
        for a in range(0, Tf, CHUNK):
            b = min(a + CHUNK, Tf); s0 = max(0, a - LEFT); e = min(Tf, b + RIGHT)
            w = self.dec.run(None, {"z": z[:, :, s0:e]})[0].reshape(-1)
            off = (a - s0) * HOP; keep = (b - a) * HOP
            if not emit(_pcm16(w[off:off + keep])):
                break


def _pcm16(x):
    x = np.clip(x, -1.0, 1.0)
    return (x * 32767.0).astype("<i2").tobytes()


BACKEND = Backend()

# ---- two preloaded gguf models: story (llama2.c-zh 15M) + chat (SmolLM-Chinese-180M) ----
import time, random, textgen
import opencc
_CC = opencc.OpenCC("s2twp")  # simplified zh -> Taiwan traditional (+phrases)
# LLM weights hosted in the model repo (not bundled in the Space), downloaded at startup.
def _llm(fn):
    from huggingface_hub import hf_hub_download
    return hf_hub_download(os.environ.get("PRIMETTS_REPO", "Luigi/PrimeTTS"), f"streaming_llm/{fn}")
STORY_GGUF = os.environ.get("STORY_GGUF") or _llm("story15m_q8.gguf")
CHAT_GGUF = os.environ.get("CHAT_GGUF") or _llm("gemma270m_it_q8.gguf")
_PHRASE_END = "。！？…；;!?\n，,、:："          # flush a clause to TTS at these
_TYPE_DELAY = float(os.environ.get("TYPE_DELAY", "0.03"))  # pacing for the template fallback
# story mode has NO user input — seed with a random opener for variety each run.
STORY_OPENERS = ["从前，", "很久以前，", "有一天，", "在一个小村庄里，", "从前有一个小男孩，",
                 "有一只小猫，", "在一片大森林里，", "很久很久以前，有一位国王，"]


def _load_warm(path, n_ctx):
    """Load a gguf and prime its compute graph (safe eval-warmup, not a streaming
    completion which corrupts state) so the first real request isn't cold."""
    from llama_cpp import Llama
    m = Llama(model_path=path, n_ctx=n_ctx,
              n_threads=int(os.environ.get("LLM_THREADS", "2")), verbose=False)
    try:
        m.eval(m.tokenize(b"\xe4\xbd\xa0\xe5\xa5\xbd"))  # "你好"
        m.reset()
    except Exception as e:
        print("[gen] warm-up skipped:", e)
    return m


class Gen:
    def __init__(self):
        self.story = None
        self.chat = None
        for name, path, ctx, attr in [("story", STORY_GGUF, 512, "story"),
                                      ("chat", CHAT_GGUF, 768, "chat")]:
            if os.path.exists(path):
                try:
                    setattr(self, attr, _load_warm(path, ctx))
                    print(f"[gen] {name} model loaded + warmed:", os.path.basename(path))
                except Exception as e:
                    print(f"[gen] {name} model unavailable ({e})")

    def has_story(self):
        return self.story is not None

    def has_chat(self):
        return self.chat is not None

    def stream_tokens(self, prompt, mode="story", min_chars=120, cap_chars=320):
        # mode="story" -> llama2.c-zh 15M, no input, random opener -> coherent tale.
        # mode="chat"  -> SmolLM-180M, user input, ChatML -> fluent response.
        # else (no model) -> instant template composer fallback.
        if mode == "story" and self.story is not None:
            seed = random.choice(STORY_OPENERS)
            yield seed
            ctx, total = seed, 0
            while total < min_chars and total < cap_chars:
                got = 0
                for o in self.story(ctx, max_tokens=140, temperature=0.9, top_p=0.9,
                                    repeat_penalty=1.1, stream=True, stop=["<s>", "</s>"]):
                    t = o["choices"][0]["text"]
                    if t:
                        yield t; ctx += t; total += len(t); got += len(t)
                        if total >= cap_chars:
                            break
                if got == 0:
                    break
        elif mode == "chat" and self.chat is not None:
            # gemma-3-270m-it is an INSTRUCT model — use its chat format with a
            # zh-TW instruction so it answers fluently in Traditional Chinese.
            p = (prompt or "").strip() or "你好"
            instr = "請用繁體中文簡短、自然地回答。" + p
            ctx = f"<start_of_turn>user\n{instr}<end_of_turn>\n<start_of_turn>model\n"
            for o in self.chat(ctx, max_tokens=200, temperature=0.7, top_p=0.9,
                               repeat_penalty=1.1, stream=True,
                               stop=["<end_of_turn>", "<eos>", "<start_of_turn>"]):
                t = o["choices"][0]["text"]
                if t:
                    yield t
        else:
            for tok in textgen.stream():
                time.sleep(_TYPE_DELAY)
                yield tok


CHAT = Gen()
app = FastAPI()


@app.get("/")
async def index():
    return HTMLResponse((HERE / "static" / "index.html").read_text(encoding="utf-8"))


VOICES = [{"sid": 0, "name": "Xinran (♀)"}, {"sid": 1, "name": "Anchen (♂)"},
          {"sid": 2, "name": "Bowen (♂)"}]


@app.get("/voices")
async def voices():
    return {"voices": VOICES}


@app.get("/healthz")
async def healthz():
    return {"backend": BACKEND.kind, "sr": SR, "chunk_ms": CHUNK * HOP // (SR // 1000)}


@app.websocket("/ws")
async def ws(sock: WebSocket):
    await sock.accept()
    loop = asyncio.get_event_loop()
    try:
        while True:
            msg = json.loads(await sock.receive_text())
            text = (msg.get("text") or "").strip()
            if not text:
                continue
            ns = float(msg.get("noise_scale", 0.667)); ls = float(msg.get("length_scale", 1.0))
            sid = int(msg.get("sid", 0))
            # Split into clauses (string-only, no g2p yet) so BOTH the g2p frontend and the
            # encoder run per-clause inside the worker: first audio depends on clause 1, not
            # the whole text. (g2pw is BERT-on-CPU — running it over all clauses upfront was
            # the real first-audio bottleneck on long text.)
            clauses = _split_clauses(text)
            await sock.send_text(json.dumps({"type": "start", "sr": SR}))
            q: asyncio.Queue = asyncio.Queue()

            def emit(pcm):  # called from the worker thread
                loop.call_soon_threadsafe(q.put_nowait, pcm)
                return True

            def work():
                for cl in clauses:
                    o = F.text_to_ids(cl)   # g2p this clause only
                    if not o["phone_ids"]:
                        continue
                    BACKEND.stream(o["phone_ids"], o["tone_ids"], o["lang_ids"], ns, ls, emit, sid=sid)
                loop.call_soon_threadsafe(q.put_nowait, None)

            fut = loop.run_in_executor(None, work)
            while True:
                pcm = await q.get()
                if pcm is None:
                    break
                await sock.send_bytes(pcm)
            await fut
            await sock.send_text(json.dumps({"type": "end"}))
    except WebSocketDisconnect:
        pass


@app.get("/has_chat")
async def has_chat():
    return {"chat": True, "story": CHAT.has_story(), "chat_llm": CHAT.has_chat(),
            "story_name": "llama2.c-zh-15M", "chat_name": "Gemma-3-270m-it"}


@app.websocket("/chat")
async def chat(sock: WebSocket):
    """STREAMING INPUT demo: LLM streams tokens (text) → phrase-buffered into the
    TTS → audio streams out. Both over one socket: JSON {type:token,t} for the
    live transcript, binary PCM for audio."""
    await sock.accept()
    loop = asyncio.get_event_loop()
    try:
        while True:
            msg = json.loads(await sock.receive_text())
            prompt = (msg.get("prompt") or "").strip()
            mode = msg.get("mode", "story")
            sid = int(msg.get("sid", 0))
            if mode == "chat" and not prompt:   # story needs no input; chat does
                continue
            await sock.send_text(json.dumps({"type": "start", "sr": SR}))
            q: asyncio.Queue = asyncio.Queue()

            def push(item):  # (kind, payload) from worker thread
                loop.call_soon_threadsafe(q.put_nowait, item)

            def synth_phrase(phrase):
                phrase = phrase.strip()
                if not phrase:
                    return
                o = F.text_to_ids(phrase)
                if not o["phone_ids"]:
                    return
                BACKEND.stream(o["phone_ids"], o["tone_ids"], o["lang_ids"], 0.667, 1.0,
                               lambda pcm: (push(("pcm", pcm)), True)[1], sid=sid)

            def work():
                buf = ""
                for tok in CHAT.stream_tokens(prompt, mode):
                    buf += tok
                    # at each clause boundary: convert to zh-TW, then display + speak it
                    while any(c in buf for c in _PHRASE_END):
                        i = min(buf.index(c) for c in _PHRASE_END if c in buf)
                        clause = _CC.convert(buf[:i + 1])   # simplified -> Taiwan traditional
                        push(("tok", clause))               # live zh-TW transcript
                        synth_phrase(clause)                # zh-TW audio
                        buf = buf[i + 1:]
                if buf.strip():
                    clause = _CC.convert(buf); push(("tok", clause)); synth_phrase(clause)
                push(None)

            fut = loop.run_in_executor(None, work)
            while True:
                item = await q.get()
                if item is None:
                    break
                kind, payload = item
                if kind == "tok":
                    await sock.send_text(json.dumps({"type": "token", "t": payload}))
                else:
                    await sock.send_bytes(payload)
            await fut
            await sock.send_text(json.dumps({"type": "end"}))
    except WebSocketDisconnect:
        pass


app.mount("/static", StaticFiles(directory=str(HERE / "static")), name="static")