import os, json, math, time, wave, shutil
from pathlib import Path
from dataclasses import dataclass
from typing import Any, Callable
os.environ["OMP_NUM_THREADS"] = "2"

import numpy as np
import onnxruntime as ort
import sentencepiece as spm
import torch
import torchaudio
import gradio as gr
from huggingface_hub import snapshot_download

SAMPLE_MODE_GREEDY = "greedy"
SAMPLE_MODE_FIXED = "fixed"
SAMPLE_MODE_FULL = "full"
EXECUTION_PROVIDER_CPU = "cpu"

MODEL_DIR = Path(os.environ.get("MOSS_MODEL_DIR", "/app/models"))
OUTPUT_DIR = Path(os.environ.get("MOSS_OUTPUT_DIR", "/tmp/moss_output"))
OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

SENTENCE_END_PUNCTUATION = set(".!?。！？；;")
CLAUSE_SPLIT_PUNCTUATION = set(",，、；;：:")
CLOSING_PUNCTUATION = set("\"'\"')]}）】》」』")
MANIFEST_CANDIDATE_RELATIVE_PATHS = (
    "browser_poc_manifest.json",
    "MOSS-TTS-Nano-100M-ONNX/browser_poc_manifest.json",
    "MOSS-TTS-Nano-ONNX-CPU/browser_poc_manifest.json",
)
MODEL_DIR_ALIAS_MAP = {
    "MOSS-TTS-Nano-ONNX-CPU": "MOSS-TTS-Nano-100M-ONNX",
    "MOSS-Audio-Tokenizer-Nano-ONNX-CPU": "MOSS-Audio-Tokenizer-Nano-ONNX",
}
DEFAULT_TTS_REPO = "OpenMOSS-Team/MOSS-TTS-Nano-100M-ONNX"
DEFAULT_CODEC_REPO = "OpenMOSS-Team/MOSS-Audio-Tokenizer-Nano-ONNX"
DEFAULT_INTER_CHUNK_PAUSE_SHORT = 0.40
DEFAULT_INTER_CHUNK_PAUSE_LONG = 0.24


def _argmax(values):
    return int(np.argmax(values))


def _normalize_sample_mode(raw, do_sample=True):
    s = str(raw or "").strip()
    if s in {SAMPLE_MODE_GREEDY, SAMPLE_MODE_FIXED, SAMPLE_MODE_FULL}:
        return s
    if not do_sample:
        return SAMPLE_MODE_GREEDY
    return SAMPLE_MODE_FIXED


def _softmax(values):
    mx = float(np.max(values))
    shifted = np.asarray(values - mx, dtype=np.float64)
    exps = np.exp(shifted)
    return exps / np.sum(exps, dtype=np.float64)


def _sample_from_scores(values, *, do_sample, temperature, top_k, top_p, rng):
    if not do_sample:
        return _argmax(values)
    scores = np.asarray(values, dtype=np.float32).copy() / float(temperature)
    if top_k > 0 and top_k < scores.shape[0]:
        threshold = float(np.sort(scores)[::-1][top_k - 1])
        scores[scores < threshold] = float("-inf")
    if top_p > 0 and top_p < 1:
        indexed = list(enumerate(scores.tolist()))
        indexed.sort(key=lambda x: x[1], reverse=True)
        sorted_scores = np.asarray([x[1] for x in indexed], dtype=np.float32)
        sorted_probs = _softmax(sorted_scores)
        remove_mask = [False] * len(indexed)
        cumulative = 0.0
        for i, p in enumerate(sorted_probs):
            cumulative += float(p)
            if cumulative > float(top_p):
                remove_mask[i] = True
        for i in range(len(remove_mask) - 1, 0, -1):
            remove_mask[i] = remove_mask[i - 1]
        if remove_mask:
            remove_mask[0] = False
        for i, rm in enumerate(remove_mask):
            if rm:
                scores[indexed[i][0]] = float("-inf")
    probs = _softmax(scores)
    rv = float(rng.random())
    for i, p in enumerate(probs):
        rv -= float(p)
        if rv <= 0:
            return int(i)
    return _argmax(scores)


def _apply_repetition_penalty(values, prev_ids, penalty):
    if not prev_ids or penalty == 1.0:
        return values
    result = values.copy()
    for tid in set(int(x) for x in prev_ids):
        if 0 <= tid < result.shape[0]:
            result[tid] = result[tid] * penalty if result[tid] < 0 else result[tid] / penalty
    return result


def _argmax_with_repetition_penalty(values, prev_set, penalty):
    best_idx, best_val = 0, float("-inf")
    apply = bool(prev_set) and penalty != 1.0
    for i, v in enumerate(values):
        s = float(v)
        if apply and i in prev_set:
            s = s * penalty if s < 0 else s / penalty
        if s > best_val:
            best_val, best_idx = s, i
    return int(best_idx)


def _sample_assistant_text_token(text_logits, manifest, gen_defaults, rng):
    cids = np.asarray([
        int(manifest["tts_config"]["audio_assistant_slot_token_id"]),
        int(manifest["tts_config"]["audio_end_token_id"]),
    ], dtype=np.int32)
    cs = text_logits[cids]
    si = _sample_from_scores(cs, do_sample=bool(gen_defaults["do_sample"]),
                             temperature=float(gen_defaults["text_temperature"]),
                             top_k=min(int(gen_defaults["text_top_k"]), int(cs.shape[0])),
                             top_p=float(gen_defaults["text_top_p"]), rng=rng)
    return int(cids[si])


def _sample_audio_token(audio_logits, prev_ids, prev_set, gen_defaults, rng):
    rp = float(gen_defaults["audio_repetition_penalty"])
    if not bool(gen_defaults["do_sample"]):
        return _argmax_with_repetition_penalty(audio_logits, prev_set, rp)
    penalized = _apply_repetition_penalty(audio_logits, prev_ids, rp)
    return _sample_from_scores(penalized, do_sample=True,
                               temperature=float(gen_defaults["audio_temperature"]),
                               top_k=int(gen_defaults["audio_top_k"]),
                               top_p=float(gen_defaults["audio_top_p"]), rng=rng)


def _flatten3d(nested):
    d0, d1, d2 = len(nested), len(nested[0]), len(nested[0][0])
    data = np.zeros((d0 * d1 * d2,), dtype=np.int32)
    off = 0
    for i in range(d0):
        for j in range(d1):
            for k in range(d2):
                data[off] = int(nested[i][j][k])
                off += 1
    return data, [d0, d1, d2]


def _flatten2d(nested):
    d0, d1 = len(nested), len(nested[0])
    data = np.zeros((d0 * d1,), dtype=np.int32)
    off = 0
    for i in range(d0):
        for j in range(d1):
            data[off] = int(nested[i][j])
            off += 1
    return data, [d0, d1]


def _extract_last_hidden(hs):
    if hs.ndim == 2:
        return hs.astype(np.float32, copy=False)
    return hs[:, -1, :].astype(np.float32, copy=False)


def _slice_channel_major_audio(audio, start=0, end=None):
    ch = int(audio.shape[1])
    total = int(audio.shape[2])
    s = max(0, int(start))
    e = total if end is None else max(s, min(int(end), total))
    return [audio[0, c, s:e].astype(np.float32, copy=False) for c in range(ch)]


def _contains_cjk(text):
    for c in str(text or ""):
        if "\u4e00" <= c <= "\u9fff" or "\u3400" <= c <= "\u4dbf" or "\u3040" <= c <= "\u30ff" or "\uac00" <= c <= "\ud7af":
            return True
    return False


def _prepare_text_for_sentence_chunking(text):
    t = str(text or "").strip()
    if not t:
        raise ValueError("Text prompt cannot be empty.")
    t = t.replace("\r", " ").replace("\n", " ")
    while "  " in t:
        t = t.replace("  ", " ")
    if _contains_cjk(t):
        if t[-1] not in SENTENCE_END_PUNCTUATION:
            t += "。"
        return t
    if t[:1].islower():
        t = t[:1].upper() + t[1:]
    if t[-1].isalnum():
        t += "."
    if len([x for x in t.split() if x]) < 5:
        t = f" {t}"
    return t


def _split_by_punct(text, punct):
    sentences, cur, i = [], [], 0
    while i < len(text):
        c = text[i]
        cur.append(c)
        if c in punct:
            la = i + 1
            while la < len(text) and text[la] in CLOSING_PUNCTUATION:
                cur.append(text[la])
                la += 1
            s = "".join(cur).strip()
            if s:
                sentences.append(s)
            cur.clear()
            while la < len(text) and text[la].isspace():
                la += 1
            i = la
            continue
        i += 1
    tail = "".join(cur).strip()
    if tail:
        sentences.append(tail)
    return sentences


def _merge_audio_channels(channels):
    if not channels:
        return np.zeros((0, 1), dtype=np.float32)
    if len(channels) == 1:
        return np.asarray(channels[0], dtype=np.float32).reshape(-1, 1)
    ml = min(int(c.shape[0]) for c in channels)
    return np.stack([np.asarray(c[:ml], dtype=np.float32) for c in channels], axis=1)


def _concat_waveforms(wfs):
    if not wfs:
        return np.zeros((0, 1), dtype=np.float32)
    ne = [w for w in wfs if w.size > 0]
    if not ne:
        return np.zeros((0, max(1, int(wfs[0].shape[1]) if wfs[0].ndim > 1 and wfs[0].shape[1] > 0 else 1)), dtype=np.float32)
    return np.concatenate(ne, axis=0)


def _write_wav(path, waveform, sr):
    p = Path(path).expanduser().resolve()
    p.parent.mkdir(parents=True, exist_ok=True)
    audio = np.asarray(waveform, dtype=np.float32)
    if audio.ndim == 1:
        audio = audio.reshape(-1, 1)
    pcm16 = np.round(np.clip(audio, -1.0, 1.0) * 32767.0).astype(np.int16)
    with wave.open(str(p), "wb") as f:
        f.setnchannels(int(pcm16.shape[1]))
        f.setsampwidth(2)
        f.setframerate(int(sr))
        f.writeframes(pcm16.tobytes())
    return p


@dataclass
class CodecStreamingSession:
    codec_meta: dict
    session: ort.InferenceSession

    def __post_init__(self):
        self.transformer_specs = list(self.codec_meta.get("streaming_decode", {}).get("transformer_offsets", []))
        self.attention_specs = list(self.codec_meta.get("streaming_decode", {}).get("attention_caches", []))
        self.state_feeds = {}
        self.reset()

    def reset(self):
        self.state_feeds = {}
        for s in self.transformer_specs:
            self.state_feeds[str(s["input_name"])] = np.zeros(tuple(s["shape"]), dtype=np.int32)
        for s in self.attention_specs:
            self.state_feeds[str(s["offset_input_name"])] = np.zeros(tuple(s["offset_shape"]), dtype=np.int32)
            self.state_feeds[str(s["cached_keys_input_name"])] = np.zeros(tuple(s["cache_shape"]), dtype=np.float32)
            self.state_feeds[str(s["cached_values_input_name"])] = np.zeros(tuple(s["cache_shape"]), dtype=np.float32)
            self.state_feeds[str(s["cached_positions_input_name"])] = np.full(tuple(s["positions_shape"]), -1, dtype=np.int32)

    def run_frames(self, frame_rows):
        if not frame_rows:
            return None
        nq = int(self.codec_meta["codec_config"]["num_quantizers"])
        fc = len(frame_rows)
        ac = np.zeros((1, fc, nq), dtype=np.int32)
        for fi, fr in enumerate(frame_rows):
            for ci in range(nq):
                ac[0, fi, ci] = int(fr[ci] if ci < len(fr) else 0)
        feeds = {"audio_codes": ac, "audio_code_lengths": np.asarray([fc], dtype=np.int32)}
        feeds.update(self.state_feeds)
        outs = self.session.run(None, feeds)
        onames = [o.name for o in self.session.get_outputs()]
        named = dict(zip(onames, outs, strict=True))
        for s in self.transformer_specs:
            self.state_feeds[str(s["input_name"])] = named[str(s["output_name"])]
        for s in self.attention_specs:
            self.state_feeds[str(s["offset_input_name"])] = named[str(s["offset_output_name"])]
            self.state_feeds[str(s["cached_keys_input_name"])] = named[str(s["cached_keys_output_name"])]
            self.state_feeds[str(s["cached_values_input_name"])] = named[str(s["cached_values_output_name"])]
            self.state_feeds[str(s["cached_positions_input_name"])] = named[str(s["cached_positions_output_name"])]
        return named["audio"], int(named["audio_lengths"].reshape(-1)[0])


def _resolve_stream_decode_frame_budget(emitted_total, sr, first_audio_at):
    if not first_audio_at or sr <= 0:
        return 1
    elapsed = max(0.0, time.perf_counter() - first_audio_at)
    lead = emitted_total / float(sr) - elapsed
    if not first_audio_at or lead < 0.20:
        return 1
    if lead < 0.55:
        return 2
    if lead < 1.10:
        return 4
    return 8


class MossTtsRuntime:
    def __init__(self, model_dir, thread_count=2, max_new_frames=375):
        self.model_dir = Path(model_dir).expanduser().resolve()
        self.thread_count = max(1, int(thread_count))
        self.manifest_path = self._find_manifest()
        self.manifest_dir = self.manifest_path.parent
        self.manifest = json.loads(self.manifest_path.read_text("utf-8"))
        if max_new_frames is not None:
            self.manifest["generation_defaults"]["max_new_frames"] = int(max_new_frames)
        self.rng = np.random.default_rng(1234)
        self.tts_meta_path = self._resolve_path(self.manifest["model_files"]["tts_meta"])
        self.codec_meta_path = self._resolve_path(self.manifest["model_files"]["codec_meta"])
        self.tts_meta = json.loads(self.tts_meta_path.read_text("utf-8"))
        self.codec_meta = json.loads(self.codec_meta_path.read_text("utf-8"))
        tok_path = str(self._resolve_path(self.manifest["model_files"].get("tokenizer_model", "tokenizer.model")))
        self.sp = spm.SentencePieceProcessor(model_file=tok_path)
        self.sessions = self._create_sessions()
        self.codec_stream = CodecStreamingSession(self.codec_meta, self.sessions["codec_decode_step"])

    def _find_manifest(self):
        for rp in MANIFEST_CANDIDATE_RELATIVE_PATHS:
            c = (self.model_dir / rp).resolve()
            if c.is_file():
                return c
        raise FileNotFoundError(f"browser_poc_manifest.json not found under {self.model_dir}")

    def _resolve_path(self, rel):
        resolved = (self.manifest_dir / Path(rel)).resolve()
        if resolved.exists():
            return resolved
        rt = str(rel).replace("\\", "/")
        for old, new in MODEL_DIR_ALIAS_MAP.items():
            frag = f"/{old}/"
            if frag in f"/{rt}/":
                rw = (self.manifest_dir / Path(rt.replace(old, new))).resolve()
                if rw.exists():
                    return rw
        return resolved

    def _session(self, p):
        opts = ort.SessionOptions()
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        opts.intra_op_num_threads = self.thread_count
        opts.inter_op_num_threads = 1
        return ort.InferenceSession(str(p), sess_options=opts, providers=["CPUExecutionProvider"])

    def _create_sessions(self):
        td = self.tts_meta_path.parent
        cd = self.codec_meta_path.parent
        sess = {
            "prefill": self._session(td / self.tts_meta["files"]["prefill"]),
            "decode": self._session(td / self.tts_meta["files"]["decode_step"]),
            "local_decoder": self._session(td / self.tts_meta["files"]["local_decoder"]),
            "codec_encode": self._session(cd / self.codec_meta["files"]["encode"]),
            "codec_decode": self._session(cd / self.codec_meta["files"]["decode_full"]),
            "codec_decode_step": self._session(cd / self.codec_meta["files"]["decode_step"]),
        }
        if self.tts_meta["files"].get("local_greedy_frame"):
            sess["local_greedy_frame"] = self._session(td / self.tts_meta["files"]["local_greedy_frame"])
        if self.tts_meta["files"].get("local_fixed_sampled_frame"):
            sess["local_fixed_sampled_frame"] = self._session(td / self.tts_meta["files"]["local_fixed_sampled_frame"])
        if self.tts_meta["files"].get("local_cached_step"):
            sess["local_cached_step"] = self._session(td / self.tts_meta["files"]["local_cached_step"])
        return sess

    def list_builtin_voices(self):
        return list(self.manifest["builtin_voices"])

    def encode_text(self, text):
        return [int(t) for t in self.sp.encode(str(text or ""), out_type=int)]

    def count_text_tokens(self, text):
        return len(self.encode_text(text))

    def _load_ref_audio(self, path):
        wf, sr = torchaudio.load(str(Path(path).expanduser().resolve()))
        wf = wf.to(torch.float32)
        tsr = int(self.codec_meta["codec_config"]["sample_rate"])
        tch = int(self.codec_meta["codec_config"]["channels"])
        if sr != tsr:
            wf = torchaudio.functional.resample(wf, sr, tsr)
        cc = int(wf.shape[0])
        if cc == tch:
            pass
        elif cc == 1 and tch > 1:
            wf = wf.repeat(tch, 1)
        elif cc > 1 and tch == 1:
            wf = wf.mean(dim=0, keepdim=True)
        else:
            raise ValueError(f"Unsupported channel conversion: {cc} -> {tch}")
        return wf.unsqueeze(0).detach().cpu().numpy().astype(np.float32, copy=False)

    def encode_ref_audio(self, path):
        wf = self._load_ref_audio(path)
        wl = int(wf.shape[-1])
        outs = self.sessions["codec_encode"].run(None, {"waveform": wf, "input_lengths": np.asarray([wl], dtype=np.int32)})
        onames = [o.name for o in self.sessions["codec_encode"].get_outputs()]
        named = dict(zip(onames, outs, strict=True))
        ac = np.asarray(named["audio_codes"], dtype=np.int32)
        cl = int(np.asarray(named["audio_code_lengths"]).reshape(-1)[0])
        nq = int(self.codec_meta["codec_config"]["num_quantizers"])
        codes = []
        for fi in range(cl):
            codes.append([int(ac[0, fi, qi]) for qi in range(nq)])
        return codes

    def resolve_prompt_codes(self, *, voice, prompt_audio_path):
        if prompt_audio_path:
            return self.encode_ref_audio(prompt_audio_path)
        v = str(voice or self.list_builtin_voices()[0]["voice"])
        row = next((x for x in self.list_builtin_voices() if x["voice"] == v), None)
        if row is None:
            raise ValueError(f"Built-in voice not found: {v}")
        return list(row["prompt_audio_codes"])

    def build_text_rows(self, token_ids):
        rw = int(self.manifest["tts_config"]["n_vq"]) + 1
        rows = []
        for tid in token_ids:
            r = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * rw
            r[0] = int(tid)
            rows.append(r)
        return rows

    def build_audio_prefix_rows(self, codes, slot_id=None):
        rw = int(self.manifest["tts_config"]["n_vq"]) + 1
        sid = int(self.manifest["tts_config"]["audio_user_slot_token_id"] if slot_id is None else slot_id)
        rows = []
        for cr in codes:
            r = [int(self.manifest["tts_config"]["audio_pad_token_id"])] * rw
            r[0] = sid
            for i in range(min(len(cr), rw - 1)):
                r[i + 1] = int(cr[i])
            rows.append(r)
        return rows

    def build_request_rows(self, codes, text_ids):
        prefix = [*self.manifest["prompt_templates"]["user_prompt_prefix_token_ids"], int(self.manifest["tts_config"]["audio_start_token_id"])]
        suffix = [int(self.manifest["tts_config"]["audio_end_token_id"]), *self.manifest["prompt_templates"]["user_prompt_after_reference_token_ids"], *text_ids, *self.manifest["prompt_templates"]["assistant_prompt_prefix_token_ids"], int(self.manifest["tts_config"]["audio_start_token_id"])]
        rows = [*self.build_text_rows(prefix), *self.build_audio_prefix_rows(codes), *self.build_text_rows(suffix)]
        return {"inputIds": rows, "attentionMask": [[1 for _ in rows]]}

    def run_local_decoder(self, gh, text_tid, frame_prefix):
        nvq = int(self.manifest["tts_config"]["n_vq"])
        apad = int(self.manifest["tts_config"]["audio_pad_token_id"])
        pp = np.full((1, nvq - 1), apad, dtype=np.int32)
        for i in range(min(len(frame_prefix), nvq - 1)):
            pp[0, i] = int(frame_prefix[i])
        outs = self.sessions["local_decoder"].run(None, {"global_hidden": gh.astype(np.float32, copy=False), "text_token_id": np.asarray([int(text_tid)], dtype=np.int32), "audio_prefix_token_ids": pp})
        on = [o.name for o in self.sessions["local_decoder"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        return nd["text_logits"].reshape(-1), nd["audio_logits"]

    def create_empty_local_past(self):
        ll = int(self.tts_meta["model_config"]["local_layers"])
        lh = int(self.tts_meta["model_config"]["local_heads"])
        lhd = int(self.tts_meta["model_config"]["local_head_dim"])
        return {n: np.zeros((1, 0, lh, lhd), dtype=np.float32) for li in range(ll) for n in (f"local_past_key_{li}", f"local_past_value_{li}")}

    def run_local_cached_step(self, gh, *, text_tid, audio_tid, ch_idx, step_type, past_vl, past):
        outs = self.sessions["local_cached_step"].run(None, {
            "global_hidden": gh.astype(np.float32, copy=False),
            "text_token_id": np.asarray([int(text_tid)], dtype=np.int32),
            "audio_token_id": np.asarray([int(audio_tid)], dtype=np.int32),
            "channel_index": np.asarray([int(ch_idx)], dtype=np.int32),
            "step_type": np.asarray([int(step_type)], dtype=np.int32),
            "past_valid_lengths": np.asarray([int(past_vl)], dtype=np.int32),
            **past,
        })
        on = [o.name for o in self.sessions["local_cached_step"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        npast = {n.replace("local_present_", "local_past_"): nd[n] for n in self.tts_meta["onnx"]["local_cached_output_names"][2:]}
        return nd["text_logits"].reshape(-1), nd["audio_logits"], npast

    def run_local_greedy_frame(self, gh, *, prev_sets, rep_penalty):
        acs = int(self.tts_meta["model_config"]["audio_codebook_sizes"][0])
        nvq = int(self.manifest["tts_config"]["n_vq"])
        rm = np.zeros((1, nvq, acs), dtype=np.int32)
        for ci, ts in enumerate(prev_sets):
            for tid in ts:
                if 0 <= tid < acs:
                    rm[0, ci, tid] = 1
        outs = self.sessions["local_greedy_frame"].run(None, {"global_hidden": gh.astype(np.float32, copy=False), "repetition_seen_mask": rm, "repetition_penalty": np.asarray([float(rep_penalty)], dtype=np.float32)})
        on = [o.name for o in self.sessions["local_greedy_frame"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        cont = bool(int(np.asarray(nd["should_continue"]).reshape(-1)[0]))
        ftids = np.asarray(nd["frame_token_ids"]).reshape(-1).astype(np.int32, copy=False).tolist()
        return cont, [int(x) for x in ftids]

    def run_local_fixed_sampled_frame(self, gh, *, prev_sets):
        acs = int(self.tts_meta["model_config"]["audio_codebook_sizes"][0])
        nvq = int(self.manifest["tts_config"]["n_vq"])
        rm = np.zeros((1, nvq, acs), dtype=np.int32)
        for ci, ts in enumerate(prev_sets):
            for tid in ts:
                if 0 <= tid < acs:
                    rm[0, ci, tid] = 1
        aru = np.asarray([min(0.99999994, max(0.0, float(self.rng.random())))], dtype=np.float32)
        au = np.asarray([[min(0.99999994, max(0.0, float(self.rng.random()))) for _ in range(nvq)]], dtype=np.float32)
        outs = self.sessions["local_fixed_sampled_frame"].run(None, {"global_hidden": gh.astype(np.float32, copy=False), "repetition_seen_mask": rm, "assistant_random_u": aru, "audio_random_u": au})
        on = [o.name for o in self.sessions["local_fixed_sampled_frame"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        ftids = np.asarray(nd["frame_token_ids"]).reshape(-1).astype(np.int32, copy=False).tolist()
        cont = bool(int(np.asarray(nd["should_continue"]).reshape(-1)[0]))
        return cont, [int(x) for x in ftids]

    def slice_audio_channel_logits(self, alogits, ci):
        pc = int(alogits.shape[-1])
        flat = alogits.reshape(-1)
        return flat[ci * pc:(ci + 1) * pc]

    def decode_full_audio(self, frames):
        if not frames:
            return [], 0
        ac, dims = _flatten3d([frames])
        outs = self.sessions["codec_decode"].run(None, {"audio_codes": ac.reshape(dims), "audio_code_lengths": np.asarray([len(frames)], dtype=np.int32)})
        on = [o.name for o in self.sessions["codec_decode"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        al = int(nd["audio_lengths"].reshape(-1)[0])
        return _slice_channel_major_audio(nd["audio"], 0, al), al

    def generate_audio_frames(self, req_rows, on_frame=None):
        gd = self.manifest["generation_defaults"]
        rw = int(self.manifest["tts_config"]["n_vq"]) + 1
        pids, pdims = _flatten3d([req_rows["inputIds"]])
        pmask, pmdims = _flatten2d(req_rows["attentionMask"])
        outs = self.sessions["prefill"].run(None, {"input_ids": pids.reshape(pdims), "attention_mask": pmask.reshape(pmdims)})
        on = [o.name for o in self.sessions["prefill"].get_outputs()]
        nd = dict(zip(on, outs, strict=True))
        gh = _extract_last_hidden(nd["global_hidden"])
        pvl = sum(int(x) for x in req_rows["attentionMask"][0])
        past = {n.replace("present_", "past_"): nd[n] for n in self.tts_meta["onnx"]["prefill_output_names"][1:]}
        gen_frames = []
        prev_by_ch = [[] for _ in range(int(self.manifest["tts_config"]["n_vq"]))]
        prev_set_by_ch = [set() for _ in range(int(self.manifest["tts_config"]["n_vq"]))]

        for si in range(int(gd["max_new_frames"])):
            frame = []
            if "local_greedy_frame" in self.sessions and not bool(gd["do_sample"]):
                cont, frame = self.run_local_greedy_frame(gh, prev_sets=prev_set_by_ch, rep_penalty=float(gd["audio_repetition_penalty"]))
                if not cont:
                    break
                for ci, st in enumerate(frame):
                    prev_by_ch[ci].append(st)
                    prev_set_by_ch[ci].add(st)
            elif "local_fixed_sampled_frame" in self.sessions and gd["sample_mode"] == SAMPLE_MODE_FIXED:
                cont, frame = self.run_local_fixed_sampled_frame(gh, prev_sets=prev_set_by_ch)
                if not cont:
                    break
                for ci, st in enumerate(frame):
                    prev_by_ch[ci].append(st)
                    prev_set_by_ch[ci].add(st)
            elif "local_cached_step" in self.sessions:
                lp = self.create_empty_local_past()
                lpvl = 0
                tl, _, lp = self.run_local_cached_step(gh, text_tid=0, audio_tid=0, ch_idx=0, step_type=0, past_vl=lpvl, past=lp)
                lpvl += 1
                ntt = _sample_assistant_text_token(tl, self.manifest, gd, self.rng)
                if ntt != int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]):
                    break
                _, alogits, lp = self.run_local_cached_step(gh, text_tid=ntt, audio_tid=0, ch_idx=0, step_type=1, past_vl=lpvl, past=lp)
                lpvl += 1
                fl = self.slice_audio_channel_logits(alogits, 0).astype(np.float32, copy=False)
                st = _sample_audio_token(fl, prev_by_ch[0], prev_set_by_ch[0], gd, self.rng)
                frame.append(st)
                prev_by_ch[0].append(st)
                prev_set_by_ch[0].add(st)
                prev = st
                for ci in range(1, int(self.manifest["tts_config"]["n_vq"])):
                    _, alogits, lp = self.run_local_cached_step(gh, text_tid=0, audio_tid=prev, ch_idx=ci - 1, step_type=2, past_vl=lpvl, past=lp)
                    lpvl += 1
                    cl = self.slice_audio_channel_logits(alogits, ci).astype(np.float32, copy=False)
                    st = _sample_audio_token(cl, prev_by_ch[ci], prev_set_by_ch[ci], gd, self.rng)
                    frame.append(st)
                    prev_by_ch[ci].append(st)
                    prev_set_by_ch[ci].add(st)
                    prev = st
            else:
                tl, _ = self.run_local_decoder(gh, 0, [])
                ntt = _sample_assistant_text_token(tl, self.manifest, gd, self.rng)
                if ntt != int(self.manifest["tts_config"]["audio_assistant_slot_token_id"]):
                    break
                for ci in range(int(self.manifest["tts_config"]["n_vq"])):
                    _, alogits = self.run_local_decoder(gh, ntt, frame)
                    cl = self.slice_audio_channel_logits(alogits, ci).astype(np.float32, copy=False)
                    st = _sample_audio_token(cl, prev_by_ch[ci], prev_set_by_ch[ci], gd, self.rng)
                    frame.append(st)
                    prev_by_ch[ci].append(st)
                    prev_set_by_ch[ci].add(st)
            gen_frames.append(frame)
            nr = np.full((1, 1, rw), int(self.manifest["tts_config"]["audio_pad_token_id"]), dtype=np.int32)
            nr[0, 0, 0] = int(self.manifest["tts_config"]["audio_assistant_slot_token_id"])
            for i, t in enumerate(frame):
                nr[0, 0, i + 1] = int(t)
            df = {"input_ids": nr, "past_valid_lengths": np.asarray([pvl], dtype=np.int32)}
            for iname in self.tts_meta["onnx"]["decode_input_names"][2:]:
                df[iname] = past[iname]
            dout = self.sessions["decode"].run(None, df)
            dn = [o.name for o in self.sessions["decode"].get_outputs()]
            dnd = dict(zip(dn, dout, strict=True))
            gh = _extract_last_hidden(dnd["global_hidden"])
            pvl += 1
            past = {n.replace("present_", "past_"): dnd[n] for n in self.tts_meta["onnx"]["decode_output_names"][1:]}
            if on_frame is not None:
                on_frame(gen_frames, si, frame)
        return gen_frames

    def decode_full_audio_safe(self, frames):
        try:
            ch_arrays, _ = self.decode_full_audio(frames)
            return _merge_audio_channels(ch_arrays)
        except Exception as exc:
            import logging
            logging.warning("full codec decode failed, falling back: %s", exc)
            self.codec_stream.reset()
            nch = int(self.codec_meta["codec_config"]["channels"])
            merged = [[] for _ in range(nch)]
            try:
                for si in range(0, len(frames), 8):
                    chunk = frames[si:si + 8]
                    dec = self.codec_stream.run_frames(chunk)
                    if dec is None:
                        continue
                    audio, al = dec
                    if al <= 0:
                        continue
                    for ci in range(nch):
                        merged[ci].append(np.asarray(audio[0, ci, :al], dtype=np.float32))
            finally:
                self.codec_stream.reset()
            return _merge_audio_channels([np.concatenate(c) if c else np.zeros((0,), dtype=np.float32) for c in merged])

    def split_text_chunks(self, text, max_tokens=75):
        t = str(text or "").strip()
        if not t:
            return []
        pieces = []
        pref = set(CLAUSE_SPLIT_PUNCTUATION) | set(SENTENCE_END_PUNCTUATION) | {" "}
        while t:
            if self.count_text_tokens(t) <= max_tokens:
                pieces.append(t)
                break
            lo, hi, best = 1, len(t), 1
            while lo <= hi:
                mid = (lo + hi) // 2
                cand = t[:mid].strip()
                if cand and self.count_text_tokens(cand) <= max_tokens:
                    best = mid
                    lo = mid + 1
                else:
                    hi = mid - 1
                if not cand:
                    lo = mid + 1
            ci = best
            pf = t[:best]
            pi = -1
            for si in range(len(pf) - 1, max(-1, len(pf) - 25), -1):
                if pf[si] in pref:
                    pi = si + 1
                    break
            if pi > 0:
                ci = pi
            piece = t[:ci].strip()
            if not piece:
                piece = t[:best].strip()
                ci = best
            pieces.append(piece)
            t = t[ci:].strip()
        return pieces if len(pieces) > 1 else [str(text or "").strip()]

    def synthesize(self, *, text, voice=None, prompt_audio_path=None, sample_mode="fixed", do_sample=True, streaming=True, max_new_frames=375):
        gd = self.manifest["generation_defaults"]
        gd["max_new_frames"] = int(max_new_frames)
        nsm = _normalize_sample_mode(sample_mode, do_sample)
        gd["sample_mode"] = nsm
        gd["do_sample"] = nsm != SAMPLE_MODE_GREEDY
        codes = self.resolve_prompt_codes(voice=voice, prompt_audio_path=prompt_audio_path)
        tid = self.encode_text(text)
        req = self.build_request_rows(codes, tid)
        if streaming:
            pending = []
            emitted = []
            emitted_total = 0
            first_at = None
            self.codec_stream.reset()

            def decode_pending(force):
                nonlocal emitted_total, first_at
                pc = len(pending)
                if pc <= 0:
                    return
                sr = int(self.codec_meta["codec_config"]["sample_rate"])
                budget = _resolve_stream_decode_frame_budget(emitted_total, sr, first_at)
                if not force and pc < max(1, budget):
                    return
                fb = pc if force else min(pc, max(1, budget))
                chunk = pending[:fb]
                del pending[:fb]
                dec = self.codec_stream.run_frames(chunk)
                if dec is None:
                    return
                audio, al = dec
                if al <= 0:
                    return
                if first_at is None:
                    first_at = time.perf_counter()
                emitted_total += al
                nch = int(self.codec_meta["codec_config"]["channels"])
                emitted.append(_merge_audio_channels([audio[0, c, :al] for c in range(nch)]))

            def on_frame(gf, si, f):
                pending.append(list(f))
                decode_pending(False)

            try:
                gf = self.generate_audio_frames(req, on_frame=on_frame)
                decode_pending(True)
            finally:
                self.codec_stream.reset()
            waveform = _concat_waveforms(emitted)
        else:
            gf = self.generate_audio_frames(req)
            waveform = self.decode_full_audio_safe(gf)

        sr = int(self.codec_meta["codec_config"]["sample_rate"])
        out_path = OUTPUT_DIR / "output.wav"
        _write_wav(out_path, waveform, sr)
        return {"audio_path": str(out_path), "sample_rate": sr, "frames": len(gf)}


def ensure_models():
    tts_dir = MODEL_DIR / "MOSS-TTS-Nano-100M-ONNX"
    codec_dir = MODEL_DIR / "MOSS-Audio-Tokenizer-Nano-ONNX"
    if not (tts_dir / "browser_poc_manifest.json").is_file():
        tts_dir.mkdir(parents=True, exist_ok=True)
        snapshot_download(DEFAULT_TTS_REPO, local_dir=str(tts_dir), local_dir_use_symlinks=False, allow_patterns=["*.onnx", "*.data", "*.json", "tokenizer.model"])
        src = tts_dir / "MOSS-TTS-Nano-100M-ONNX"
        if src.is_dir():
            for f in src.iterdir():
                dst = tts_dir / f.name
                if not dst.exists():
                    shutil.move(str(f), str(dst))
    if not (codec_dir / "codec_browser_onnx_meta.json").is_file():
        codec_dir.mkdir(parents=True, exist_ok=True)
        snapshot_download(DEFAULT_CODEC_REPO, local_dir=str(codec_dir), local_dir_use_symlinks=False, allow_patterns=["*.onnx", "*.data", "*.json"])
        src = codec_dir / "MOSS-Audio-Tokenizer-Nano-ONNX"
        if src.is_dir():
            for f in src.iterdir():
                dst = codec_dir / f.name
                if not dst.exists():
                    shutil.move(str(f), str(dst))


runtime = None


def get_runtime():
    global runtime
    if runtime is not None:
        return runtime
    ensure_models()
    runtime = MossTtsRuntime(MODEL_DIR, thread_count=2, max_new_frames=375)
    return runtime


def synthesize_gradio(text, voice, audio_path, sample_mode, max_frames):
    rt = get_runtime()
    t0 = time.time()
    result = rt.synthesize(
        text=text,
        voice=voice if not audio_path else None,
        prompt_audio_path=audio_path if audio_path else None,
        sample_mode=sample_mode,
        do_sample=(sample_mode != "greedy"),
        streaming=True,
        max_new_frames=int(max_frames),
    )
    elapsed = time.time() - t0
    return result["audio_path"], f"Done in {elapsed:.1f}s | {result['sample_rate']}Hz | {result['frames']} frames"


VOICES = ["Junhao", "Zhiming", "Weiguo", "Xiaoyu", "Yuewen", "Lingyu", "Trump", "Ava", "Bella", "Adam", "Nathan", "Soyo", "Saki", "Mortis", "Umiri", "Mei", "Anon", "Arisa"]

with gr.Blocks(title="MOSS-TTS-Nano ONNX") as demo:
    gr.Markdown("# MOSS-TTS-Nano-100M-ONNX\nCPU-only TTS with voice cloning. First run downloads ~730MB model.")
    with gr.Row():
        with gr.Column():
            text_in = gr.Textbox(label="Text", value="Hello, welcome to MOSS TTS Nano.", lines=3)
            with gr.Row():
                voice_in = gr.Dropdown(choices=VOICES, value="Junhao", label="Voice (overridden by ref audio)")
                ref_audio = gr.Audio(label="Reference Audio (optional, for voice cloning)", type="filepath")
            with gr.Row():
                sample_mode = gr.Dropdown(choices=["fixed", "greedy", "full"], value="fixed", label="Sample Mode")
                max_frames = gr.Slider(16, 750, value=375, step=1, label="Max Frames")
            btn = gr.Button("Synthesize", variant="primary")
        with gr.Column():
            audio_out = gr.Audio(label="Generated Audio", type="filepath")
            info_out = gr.Textbox(label="Info")
    btn.click(fn=synthesize_gradio, inputs=[text_in, voice_in, ref_audio, sample_mode, max_frames], outputs=[audio_out, info_out])

if __name__ == "__main__":
    get_runtime()
    demo.launch()