"""Maya 2 generation for one request at a time: the official engine's per-piece loop (maya3/engine.py) on MLX.

    from maya2_mlx import Maya2
    tts = Maya2()                 # models/maya2-8bit: dynamic 8-bit backbone, fp16 depth decoder, fp32 Mimi
    wav = tts.generate("Hello there.", lang="en", voice="Anirban", seed=7)    # float32, 24 kHz, in [-1, 1]

Same request semantics as Maya3TTS.generate (api.py:224-230 -> engine.py): split into ~20 s pieces, one fresh
backbone context per piece, controlled stop, duration cap and loop guard, 0.15 s gap between pieces, int16 PCM / 32767.
"""
import math
import os
import random
import time

import mlx.core as mx
import numpy as np

from . import philox
from . import quant as Q
from . import weights as W
from .backbone import Backbone
from .depth import Depth
from .mimi import Mimi
from .prompt import Packer, VoiceRegistry, expected_frames, load_named_voices, load_prior, split_text

PORT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
SR = 24000                     # engine.py:42
FRAME = 1920                   # engine.py:43
AEOS = 2048                    # tts/tokens.py:23
REASONS = {1: "aeos", 2: "aeos_argmax", 3: "aeos_tau", 4: "aeos_hazard"}     # engine.py:46

# EngineConfig defaults (engine.py:49-79) that the single-request loop reads
CB0_TEMP, CB0_TOP_K, CB0_TOP_P, CB0_REP, CB0_REP_WIN = 0.9, 64, 1.0, 0.8, 12
DEPTH_TEMP, DEPTH_TOP_K = 0.7, 64
HARD_CAP, CAPACITY, SEED = 1500, 2048, 20261024
GAP_S, CHUNK_S = 0.15, 20.0
WARMUP_PIECES = 2              # Engine.warmup runs two one-piece requests before any user request (engine.py:243-249)


class GenerationError(RuntimeError):
    """api.py:31-34"""

    def __init__(self, code, msg):
        super().__init__(f"{code}: {msg}")
        self.code = code


def check_sampling(sp):
    """engine.py:110-128"""
    if not sp:
        return {}
    allowed = {"seed", "temperature", "top_k", "top_p", "repetition_penalty"}
    bad = set(sp) - allowed
    if bad:
        raise ValueError(f"unknown sampling keys {sorted(bad)}; allowed {sorted(allowed)}")
    out = {k: v for k, v in sp.items() if v is not None}
    if "temperature" in out and not 0.05 <= float(out["temperature"]) <= 2.0:
        raise ValueError("temperature must be in [0.05, 2.0]")
    if "top_k" in out and not 1 <= int(out["top_k"]) <= 2049:
        raise ValueError("top_k must be in [1, 2049]")
    if "top_p" in out and not 0.0 < float(out["top_p"]) <= 1.0:
        raise ValueError("top_p must be in (0, 1]")
    if "repetition_penalty" in out and not 0.0 <= float(out["repetition_penalty"]) <= 5.0:
        raise ValueError("repetition_penalty must be in [0, 5]")
    if "seed" in out:
        out["seed"] = int(out["seed"])
    return out


class Maya2:
    def __init__(self, models_dir=W.DEFAULT_DIR, strict=False):
        """strict=False: production precision (backbone bf16, depth fp16, Mimi fp32 without TF32).
        strict=True: fp32 everywhere (the parity gates)."""
        from maya3 import guards as G                              # the official file, unchanged
        self.G = G
        t0 = time.perf_counter()
        meta = W.load_meta(models_dir)
        self.meta = meta
        self.models_dir = models_dir
        self.bdt = mx.float32 if strict else mx.bfloat16              # engine.py:158 (model dtype)
        self.ddt = mx.float32 if strict else mx.float16               # engine.py:192, depth_dtype "fp16"
        self.backbone = Backbone(meta, self.bdt, models_dir, capacity=CAPACITY)
        self.recipe = Q.load_recipe(models_dir)
        if strict and self.recipe:
            raise ValueError("strict (fp32) runs use the bf16 release export, not quantized weights")
        aud = W.load(["audio_in.safetensors"], models_dir)
        table = Q.cast(Q.take(aud, "table.weight", self.recipe), self.bdt)
        fbw = aud["norm.weight"].astype(self.bdt)
        self.depth = Depth(meta, (table, fbw), self.ddt, DEPTH_TEMP, DEPTH_TOP_K, models_dir)
        self.depth.set_speaker(None)                                   # named / plain: the null speaker memory
        self.mimi = Mimi(meta, models_dir)
        self.anchor = self._anchor(W.load(["anchor.safetensors"], models_dir))
        self.packer = Packer(os.path.join(models_dir, "tokenizer.json"), meta["token_ids"])
        self.prior = load_prior(os.path.join(PORT, "configs", "text_frames_prior.json"))
        self.dur = G.DurationPrior()
        self.gcfg = G.GuardConfig()
        self.registry = VoiceRegistry()
        load_named_voices(self.registry, os.path.join(models_dir, "voices.json"))
        self._seed_rng = random.Random(SEED)                          # engine.py:224-225
        for _ in range(WARMUP_PIECES):
            self._seed_rng.randrange(1, 2 ** 31 - 1)
        self.load_s = time.perf_counter() - t0
        self.last_info = None

    # ------------------------------------------------------------------ constants
    def _anchor(self, w):
        """The speaker anchor of a row with no speaker embedding (anchors.py:86-109): fc2(gelu(fc1(null_vec))),
        Linear = fp32 accumulation + bias, one rounding; GELU in fp32, one rounding. Model dtype."""
        dt = self.bdt
        x = w["null_vec"].astype(dt)[None]
        h = (x.astype(mx.float32) @ w["fc1.weight"].astype(dt).astype(mx.float32).T
             + w["fc1.bias"].astype(dt).astype(mx.float32)).astype(dt)
        hf = h.astype(mx.float32)
        h = (hf * 0.5 * (1.0 + mx.erf(hf * math.sqrt(0.5)))).astype(dt)
        a = (h.astype(mx.float32) @ w["fc2.weight"].astype(dt).astype(mx.float32).T
             + w["fc2.bias"].astype(dt).astype(mx.float32)).astype(dt)
        mx.eval(a)
        return a[0]

    @property
    def voices(self):
        return self.registry.ids()

    # ------------------------------------------------------------------ public
    def generate(self, text, lang="en", voice=None, max_frames=None, return_info=False, **kw):
        """Synthesise text -> float32 waveform (24 kHz), as Maya3TTS.generate (api.py:224-230).
        kw: normalize, seed, temperature, top_k, top_p, repetition_penalty. (The official positional arguments
        ref_audio, ref_text, ref_lang and mode belong to cloning, which is not ported; pass the rest by keyword.)"""
        info = {}
        chunks = list(self.stream(text, lang, voice, max_frames=max_frames, info=info, **kw))
        wav = (np.concatenate(chunks).astype(np.float32) / 32767.0) if chunks else np.zeros(0, np.float32)
        return (wav, info) if return_info else wav

    def stream(self, text, lang="en", voice=None, max_frames=None, normalize=False, info=None, seed=None,
               temperature=None, top_k=None, top_p=None, repetition_penalty=None):
        """Yield int16 PCM chunks (24 kHz) as they are generated (api.py:196-222); `info` (a dict) receives the end
        payload. Errors while preparing a piece raise GenerationError("invalid_request", ...) as the engine reports
        them (engine.py:462-468, api.py:219-220)."""
        v = self.registry.get(voice or "plain")                         # api.py:203 (before the sampling check)
        if normalize:                                                   # api.py:188-194
            from maya3.text import normalise_text
            text = [normalise_text(t, lang) for t in text] if isinstance(text, (list, tuple)) \
                else normalise_text(text, lang)
        sampling = check_sampling({"seed": seed, "temperature": temperature, "top_k": top_k, "top_p": top_p,
                                   "repetition_penalty": repetition_penalty})                  # engine.py:293
        if not isinstance(text, (list, tuple)):                         # engine.py:296-299
            if not isinstance(text, str) or not text.strip():
                raise ValueError("text must be a non-empty string")
            text = text.strip()
            pieces = split_text(text, lang, self.prior, CHUNK_S) if CHUNK_S > 0 else [text]
        else:
            pieces = list(text)
        req = {"frames": 0, "stops": [], "t0": time.perf_counter(), "t_first": None}
        k = 0
        while True:
            try:
                plan = self._prepare(pieces[k], lang, v, max_frames, req)
            except Exception as exc:                                    # engine.py:464-468
                raise GenerationError("invalid_request", f"{type(exc).__name__}: {exc}") from exc
            for pcm in self._piece(plan, k, sampling, max_frames, req):
                yield pcm
            k += 1
            if k < len(pieces) and not (max_frames and req["frames"] >= max_frames):     # engine.py:641-645
                yield np.zeros(int(GAP_S * SR), np.int16)
            else:
                break
        end = {"reason": req["stops"][-1]["reason"], "frames": req["frames"], "audio_s": req["frames"] * FRAME / SR,
               "ttfa_s": (req["t_first"] - req["t0"]) if req["t_first"] else None,
               "wall_s": time.perf_counter() - req["t0"], "pieces": len(pieces), "stops": req["stops"]}
        if info is not None:
            info.update(end)
        self.last_info = end

    # ------------------------------------------------------------------ one piece
    def _plan(self, text, lang, max_frames, req):
        """engine.py:394-418 + the capacity rule of engine.py:469-476."""
        g = self.gcfg
        ef, er = expected_frames(self.prior, lang, len(text))
        loose = min(HARD_CAP, int(2.0 * (ef + er) + 40))
        cap = loose
        if g.enabled:
            cap = min(cap, max(int(ef * 0.5) + 8, self.dur.cap_frames(text, lang, g.dur_k, g.dur_margin_s)))
        if max_frames:
            cap = min(cap, max(1, int(max_frames) - req["frames"]))
        minimum = max(4, int(0.5 * ef))
        if g.enabled and g.loop and g.aware_v2:
            suspend = self.G.suspend_windows_v2(text, lang, self.dur)
        else:
            suspend = self.G.tag_windows(text, lang, self.dur, g.tag_hold_s) \
                if (g.enabled and g.loop and g.tag_aware and "[" in text) else ()
        return cap, loose, minimum, suspend

    def _prompt(self, voice, text, lang):
        """Token ids and the composed prompt embeddings (s2_serve_controls.py:197-223, tts_model_v2.py:128-170):
        text embeddings, plus the null speaker anchor at the <|audio|> sentinel (the last position)."""
        ids = self.packer.prompt_ids(voice.render(text), lang)
        emb = self.backbone.embed_tokens(ids)
        last = emb[-1] + self.anchor                                     # index_put accumulate: one add
        emb = mx.concatenate([emb[:-1], last[None]], 0)
        return ids, emb

    def _cb0(self, logits, hist, hz, can, sp, aeos_bias, noise):
        """Engine._cb0 (engine.py:652-694) for one seat. logits fp32 [2049]; hist int [12] (-1 empty);
        hz fp32 scalar; can bool. Returns (stop code, cb0, new hz) as mx arrays."""
        V = logits.shape[0]
        idx = mx.arange(V)
        temp = float(sp.get("temperature") or CB0_TEMP)
        top_k = int(sp.get("top_k") or CB0_TOP_K)
        top_p = float(sp.get("top_p") or CB0_TOP_P)
        rep = float(sp["repetition_penalty"] if sp.get("repetition_penalty") is not None else CB0_REP)
        logits = logits / mx.array(temp, mx.float32)
        logits = logits + mx.where(idx == AEOS, mx.array(aeos_bias, mx.float32), mx.array(0.0, mx.float32))
        if CB0_REP_WIN > 0:
            h = mx.array(hist, mx.int32)
            pen = mx.zeros((V,), mx.float32).at[mx.maximum(h, 0)].add((h >= 0).astype(mx.float32))
            pen = mx.where(idx == AEOS, 0.0, pen)
            logits = logits - mx.array(rep, mx.float32) * (pen > 0).astype(mx.float32)
        srt = mx.sort(logits)                                            # ascending: desc index i = V-1-i
        kth = srt[V - 1 - min(max(top_k - 1, 0), V - 1)]
        logits = mx.where(logits < kth, -mx.inf, logits)
        if top_p < 1.0:
            p = mx.softmax(logits, precise=True)
            si = mx.argsort(-p)
            sp_ = p[si]
            drop_sorted = (mx.cumsum(sp_) - sp_) > top_p
            drop = mx.zeros((V,), mx.int32).at[si].add(drop_sorted.astype(mx.int32)) > 0   # scatter back by si
            logits = mx.where(drop, -mx.inf, logits)
        probs = mx.softmax(logits, precise=True)
        pa = probs[AEOS]
        arg = mx.argmax(probs)
        hzv = mx.array(hz, mx.float32)
        if can:
            hzv = 1 - (1 - hzv) * (1 - pa)
        code = mx.array(0)
        if can:
            code = mx.where(hzv >= 0.6, 4, code)
            code = mx.where(pa >= 0.35, 3, code)
            code = mx.where(arg == AEOS, 2, code)
        q = mx.where(idx == AEOS, 0.0, probs)
        q = q / mx.maximum(q.sum(), 1e-12)
        cb0 = mx.minimum(mx.argmax(q / noise), AEOS - 1).astype(mx.int32)
        return code, cb0, hzv

    def _prepare(self, text, lang, voice, max_frames, req):
        """_plan (engine.py:400-418: the prompt is built first) and the capacity rule of _admit (:469-476)."""
        ids, emb = self._prompt(voice, text, lang)
        cap, loose, minimum, suspend = self._plan(text, lang, max_frames, req)
        n = len(ids)
        if n + cap + 2 > CAPACITY:
            cap = CAPACITY - n - 2
            if cap < 8:
                raise ValueError(f"prompt of {n} tokens leaves no room in KV capacity {CAPACITY}")
        return emb, cap, loose, minimum, suspend

    def _piece(self, plan, k, sp, max_frames, req):
        g = self.gcfg
        emb, cap, loose, minimum, suspend = plan
        if "seed" in sp and sp["seed"] is not None:                      # engine.py:430-437
            seed = (int(sp["seed"]) * 1000003 + k * 7919) % (2 ** 31 - 1)
        else:
            seed = self._seed_rng.randrange(1, 2 ** 31 - 1)
        row_seed = (seed * 69069 + 1) % (2 ** 31 - 1)
        self.backbone.reset()
        self.mimi.reset()
        hist = np.full(CB0_REP_WIN, -1, np.int64)                        # cb0 presence-penalty window
        ghist = np.full(g.hist, -1, np.int64)                            # loop-guard cb0 tail (_Seat.hist)
        hz = np.float32(0.0)
        aeos_bias = 0.0
        frames, boost_at, n_boost, ctr = 0, None, 0, 0
        h = self.backbone.forward(emb)[-1]
        reason = None
        while reason is None:
            noise = mx.array(philox.cb0_noise(seed, ctr))
            uni = mx.array(philox.depth_uniform(row_seed, ctr))
            code, cb0, hz_new = self._cb0(self.backbone.cb0_logits(h), hist, hz, frames >= minimum, sp, aeos_bias,
                                          noise)
            codes, fb = self.depth.frame(h, cb0, uni)
            pcm = self.mimi.decode_frame(codes)
            mx.eval(code, cb0, hz_new, codes, fb, pcm)
            ctr += 1
            c0 = int(cb0.item())
            hist = np.concatenate([hist[1:], [c0]])
            hz = np.float32(hz_new.item())
            stop = int(code.item())
            if stop:                                                     # engine.py:576-578: not emitted
                reason = REASONS[stop]
                break
            frames += 1
            req["frames"] += 1
            if req["t_first"] is None:
                req["t_first"] = time.perf_counter()
            pcm16 = np.rint(np.clip(np.array(pcm, np.float32), -1, 1) * np.float32(32767)).astype(np.int16)
            yield pcm16
            ghist[:-1] = ghist[1:]
            ghist[-1] = c0
            if frames >= cap:                                            # engine.py:586-593
                if max_frames and req["frames"] >= max_frames:
                    reason = "max_frames"
                elif g.enabled and cap < loose:
                    reason = "cap_guard"
                else:
                    reason = "cap"
                break
            if g.enabled and g.loop and frames >= g.start and frames % g.every == 0 and \
                    not any(a <= frames <= b for a, b in suspend):       # engine.py:594-617
                hit = self.G.loop_detect_batch(ghist[None], np.array([frames]), g)[0]
                if hit is not None and g.mode == "boost":
                    if boost_at is None:
                        boost_at = frames
                        n_boost += 1
                        aeos_bias = g.boost
                    if frames - boost_at >= g.backstop:
                        reason = "loop_backstop"
                        break
                elif hit is not None:
                    reason = "loop_" + hit.split(":")[0]
                    break
                elif boost_at is not None:
                    aeos_bias = 0.0
                    boost_at = None
            h = self.backbone.forward(fb[None])[0]
        # engine.py:639, plus the piece's seed (port addition: the official stream is not reported)
        req["stops"].append({"reason": reason, "frames": frames, "cap": cap, "n_boost": n_boost, "seed": seed})