"""Streaming TTS with Flash-DPO (Dynamic Accent, CosyVoice3 engine): the tagged reply is cut into chunks
(a short first chunk for latency, then sentence by sentence) and each chunk is synthesised as soon as the previous is done."""
import os, sys, re, time
import numpy as np
REL = os.environ.get("DA_REL", "/home/dev/fi/reltest")
sys.path.insert(0, REL)
os.environ.setdefault("CV3_DIR", "/home/dev/models/Fun-CosyVoice3-0.5B-2512")
os.environ.setdefault("DA_PROMPTS", REL + "/prompts")
TAG = re.compile(r"<(eng|ibo|yor|hau|pcm)>(.*?)</\1>", re.S)
VOICE = {"anchor_male": "anchor", "waxal_ibo_5_female": "waxal_ibo_5", "waxal_yor_5_female": "waxal_yor_5"}

def chunk_tagged(tagged, first_max_words=6):
    """-> list of tagged chunk strings. Splits at sentence ends (. ! ?), keeps tags, first chunk short (also cut at a comma/greeting)."""
    units, cur = [], []
    for l, t in TAG.findall(tagged):
        parts = re.split(r"(?<=[.!?])\s+", t.strip())
        for k, s in enumerate(parts):
            if not s: continue
            cur.append((l, s))
            if re.search(r"[.!?]$", s):
                units.append(cur); cur = []
    if cur: units.append(cur)
    if not units: return []
    first = units[0]
    n = sum(len(s.split()) for _, s in first)
    if n > first_max_words:                       # fast first chunk: cut the first sentence at a comma / tag boundary / conjunction
        if len(first) > 1 and len(first[0][1].split()) <= first_max_words:
            units[0] = [first[0]]; units.insert(1, first[1:])
        else:
            l, s = first[0]
            m = re.search(r"[,;:]\s", s)
            if not (m and 1 <= len(s[:m.start()].split()) <= first_max_words):
                m = None
                ws = s.split()
                for k in range(3, min(first_max_words, len(ws) - 3)):
                    if ws[k].lower() in ("and", "but", "so", "because", "then", "or"):
                        m = re.search(r"\s" + re.escape(ws[k]) + r"\s", s); break
            if m:
                cut = m.end() if s[m.start()] in ",;:" else m.start()
                units[0] = [(l, s[:cut].strip())]; units.insert(1, [(l, s[cut:].strip())] + first[1:])
    # hard caps: first chunk <= 10 words, later chunks <= 22 words (split at a comma if possible, else at a word boundary)
    def split_unit(u, cap):
        out, cur, n = [], [], 0
        for l, s in u:
            ws = s.split()
            while n + len(ws) > cap:
                k = cap - n
                sub = " ".join(ws[:k]); c = sub.rfind(",")
                if c > len(sub) // 3: k = len(sub[:c + 1].split())
                if k > 0: cur.append((l, " ".join(ws[:k])))
                out.append(cur); cur, n, ws = [], 0, ws[k:]
            if ws: cur.append((l, " ".join(ws))); n += len(ws)
        if cur: out.append(cur)
        return out
    capped = []
    for i, u in enumerate(units):
        capped += split_unit(u, 10 if i == 0 else 22)
    units = [u for u in capped if u]
    return [" ".join(f"<{l}>{s}</{l}>" for l, s in u) for u in units]

class TTS:
    def __init__(self, voice="anchor_male", lm=None, diac=None, tail_pad=4):
        import flash_lm
        from frontend import FrontEnd
        lm = lm or f"{REL}/checkpoints/flash/flash_dpo/best.pt"
        diac = diac or f"{REL}/checkpoints/flash/diacritics_byt5"
        self.S = flash_lm.Synth(lm)
        self.fe = FrontEnd(diac if os.path.isdir(diac) else None)
        self.sr = self.S.cosy.sample_rate
        self.voice = VOICE.get(voice, voice); self.tail_pad = tail_pad
        self.synth("<eng>Warm up.</eng>")          # compile/cache voice prompt

    def synth(self, tagged, seed=0, voice=None):
        import torch
        v = VOICE.get(voice, voice) if voice else self.voice
        t0 = time.perf_counter()
        # diacritic restoration only for ibo/yor spans written without any tone marks (the LLM usually writes them)
        need = any(l in ("ibo", "yor") and not re.search(r"[\u0300-\u036f\u1ea0-\u1ef9ịụọẹṣńḿ]", __import__("unicodedata").normalize("NFD", s) + s) for l, s in TAG.findall(tagged))
        t = self.fe(tagged, "codeswitch", "eng") if need else tagged
        t1 = time.perf_counter()
        toks = list(self.S.tokens(t, v, "normal", seed)[0]); torch.cuda.synchronize()
        t2 = time.perf_counter()
        if self.tail_pad: toks += [1] * self.tail_pad
        w = np.asarray(self.S.wav(toks, v), dtype=np.float32)
        t3 = time.perf_counter()
        self.last = dict(fe=round(t1 - t0, 3), lm=round(t2 - t1, 3), flow_voc=round(t3 - t2, 3), n_tok=len(toks), audio_s=round(len(w) / self.sr, 2))
        return w

    def stream(self, tagged, seed=0, voice=None):
        """yields (chunk_text, wav, synth_seconds)"""
        for i, c in enumerate(chunk_tagged(tagged)):
            t = time.perf_counter(); w = self.synth(c, seed + 101 * i, voice)
            yield c, w, time.perf_counter() - t