"""
HF Inference Endpoint handler — Hypa Orpheus TTS + Voice Cloning (Step-III merged 16-bit).

BACKWARD COMPATIBLE with the legacy hypaai_orpheus_v5 API: legacy clients work by
changing only the endpoint URL. Legacy schema honored:
  data: inputs, clone, clone_on_the_fly, enroll_user, cloning_features, enrollments
  parameters: voice (default "Eniola"), temperature, top_p, max_new_tokens, repetition_penalty
Legacy output honored: audio_b64 = base64 WAV/RIFF PCM_16 @24kHz, audio_sample =
raw float32 mono waveform, sample_rate, input_ids_len, gen_ids_len.

NEW capabilities (Step-III model) via parameters:
  task ("tts"|"vc"), mode ("vanilla"|"translate"), language, method ("m1"|"m2"),
  reference_text + reference_audio (base64), top_k.

Serving notes vs legacy handler (deliberate changes):
  - Prompts reach vLLM as token ids (legacy decoded to a string and re-tokenized,
    risking a double-BOS and mangled audio tokens in cloning prompts).
  - dtype="bfloat16" forced (legacy inherited config torch_dtype; a bf16-trained
    model served in fp16 is the leading suspect for endpoint-only audio artifacts).
  - Reference codes are frame-deduped to match Step-III training data.
"""

import io
import os
import base64
import tempfile
import traceback

import numpy as np
import torch
import soundfile as sf
import librosa

from transformers import AutoTokenizer
from snac import SNAC
from vllm import LLM, SamplingParams


class EndpointHandler:
    TOKENISER_LEN   = 128256
    START_OF_TEXT   = 128000
    END_OF_TEXT     = 128009
    START_OF_SPEECH = TOKENISER_LEN + 1   # 128257
    END_OF_SPEECH   = TOKENISER_LEN + 2   # 128258
    START_OF_HUMAN  = TOKENISER_LEN + 3   # 128259
    END_OF_HUMAN    = TOKENISER_LEN + 4   # 128260
    START_OF_AI     = TOKENISER_LEN + 5   # 128261
    END_OF_AI       = TOKENISER_LEN + 6   # 128262
    AUDIO_OFFSET    = 128266

    MAX_MODEL_LEN   = 4096
    MAX_REF_SECONDS = 30
    SNAC_SR         = 24000

    LANG_DISPLAY = {
        "en": "English", "es": "Spanish", "fr": "French", "ha": "Hausa",
        "yo": "Yoruba", "sw": "Swahili", "ar": "Arabic", "pt": "Portuguese",
        "ann": "Annang", "ebi": "Ebira", "efi": "Efik", "ego": "Eggon",
        "urh": "Urhobo", "ibb": "Ibibio", "idm": "Idoma", "igl": "Igala",
        "ig": "Igbo", "nup": "Nupe", "tiv": "Tiv", "pg": "Pidgin",
    }

    # ------------------------------------------------------------------ init
    def __init__(self, path=""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(self.device).eval()
        self.model = LLM(
            path,
            max_model_len=self.MAX_MODEL_LEN,
            gpu_memory_utilization=0.75,
            dtype="bfloat16",            # match training numerics (see docstring)
        )
        self.tokenizer = AutoTokenizer.from_pretrained(path)

    # ------------------------------------------------------- text encoding
    def _lang_display(self, x):
        if x is None:
            return None
        k = str(x).strip().lower()
        return self.LANG_DISPLAY.get(k, k.capitalize() if k else None)

    def _encode_text_ids(self, text, speaker=None, lang_tag=None):
        """Training-identical text content: '{spk} - {Lang}: {text}' variants.
        Returns bare content ids WITHOUT specials (block adds them)."""
        text = "" if text is None else str(text).strip()
        spk = speaker if (speaker and str(speaker).strip().lower() not in ("", "random", "none")) else None
        if spk and lang_tag:
            prompt = f"{spk} - {lang_tag}: {text}"
        elif spk:
            prompt = f"{spk}: {text}"
        elif lang_tag:
            prompt = f"{lang_tag}: {text}"
        else:
            prompt = text
        return self.tokenizer.encode(prompt, add_special_tokens=False)

    def _text_block(self, content_ids, with_bos=True):
        """[SOH] (+BOS) content [EOT] [EOH] — equals legacy format_text_block and
        training's [SOH]+encode(add_bos)+[EOT]+[EOH] (BOS == START_OF_TEXT)."""
        bos = [self.START_OF_TEXT] if with_bos else []
        return [self.START_OF_HUMAN] + bos + list(content_ids) + [self.END_OF_TEXT, self.END_OF_HUMAN]

    def _audio_block(self, codes):
        return [self.START_OF_AI, self.START_OF_SPEECH] + list(codes) + \
               [self.END_OF_SPEECH, self.END_OF_AI]

    def _open_speech(self):
        return [self.START_OF_AI, self.START_OF_SPEECH]

    # ------------------------------------------------------ audio encoding
    def _b64_to_wave(self, b64_str):
        raw = base64.b64decode(b64_str)
        if not raw:
            raise ValueError("reference audio is empty.")
        try:
            arr, sr = sf.read(io.BytesIO(raw), dtype="float32")
        except Exception:
            tmp = None
            try:
                with tempfile.NamedTemporaryFile(delete=False, suffix=".audio") as f:
                    f.write(raw)
                    tmp = f.name
                arr, sr = librosa.load(tmp, sr=None, mono=False)
                arr = np.asarray(arr, dtype=np.float32)
                if arr.ndim > 1:
                    arr = arr.T
            finally:
                if tmp and os.path.exists(tmp):
                    os.remove(tmp)
        if arr.ndim > 1:
            arr = arr.mean(axis=1)
        if arr.size == 0 or not np.isfinite(arr).all():
            raise ValueError("Reference audio is empty or contains invalid samples.")
        if sr != self.SNAC_SR:
            arr = librosa.resample(arr.astype(np.float32), orig_sr=sr, target_sr=self.SNAC_SR)
        dur = len(arr) / self.SNAC_SR
        if dur > self.MAX_REF_SECONDS:
            raise ValueError(f"Reference audio is {dur:.1f}s; max is {self.MAX_REF_SECONDS}s.")
        return arr.astype(np.float32)

    @torch.inference_mode()
    def _audio_to_codes(self, arr):
        wav = torch.from_numpy(arr).to(self.device)[None, None]
        codes = self.snac_model.encode(wav)
        c0, c1, c2 = codes[0][0].tolist(), codes[1][0].tolist(), codes[2][0].tolist()
        n = min(len(c0), len(c1) // 2, len(c2) // 4)
        out = []
        for i in range(n):
            out += [
                c0[i]         + self.AUDIO_OFFSET,
                c1[2 * i]     + self.AUDIO_OFFSET + 4096,
                c2[4 * i]     + self.AUDIO_OFFSET + 2 * 4096,
                c2[4 * i + 1] + self.AUDIO_OFFSET + 3 * 4096,
                c1[2 * i + 1] + self.AUDIO_OFFSET + 4 * 4096,
                c2[4 * i + 2] + self.AUDIO_OFFSET + 5 * 4096,
                c2[4 * i + 3] + self.AUDIO_OFFSET + 6 * 4096,
            ]
        return out

    @staticmethod
    def _dedup_frames(codes):
        if not codes:
            return codes
        codes = list(codes)[: (len(codes) // 7) * 7]
        if len(codes) < 7:
            return codes
        result = codes[:7]
        for i in range(7, len(codes), 7):
            if codes[i] != result[-7]:
                result.extend(codes[i:i + 7])
        return result

    # ------------------------------------------------------ legacy enrollment
    def enroll_user(self, enrollment_pairs):
        """Legacy-format enrollment: torch-serialized {text_ids tensor, audio_codes list}.
        Previously issued cloning_features blobs remain loadable."""
        enrollment_data = []
        for text, base64_audio in enrollment_pairs:
            text_ids = self.tokenizer.encode(text, return_tensors="pt",
                                             add_special_tokens=False).cpu()
            audio_codes = self._dedup_frames(self._audio_to_codes(self._b64_to_wave(base64_audio)))
            enrollment_data.append({"text_ids": text_ids, "audio_codes": audio_codes})
        buffer = io.BytesIO()
        torch.save(enrollment_data, buffer)
        buffer.seek(0)
        return base64.b64encode(buffer.read()).decode("utf-8")

    # --------------------------------------------------------- generation
    def _generate(self, prompt_ids, gp):
        sampling = SamplingParams(
            temperature        = gp["temperature"],
            top_p              = gp["top_p"],
            top_k              = gp["top_k"],
            max_tokens         = gp["max_new_tokens"],
            repetition_penalty = gp["repetition_penalty"],
            stop_token_ids     = [self.END_OF_SPEECH, self.END_OF_AI],
            detokenize         = False,
        )
        outputs = self.model.generate({"prompt_token_ids": prompt_ids}, sampling)
        return list(outputs[0].outputs[0].token_ids)

    # ----------------------------------------------------------- decoding
    @torch.inference_mode()
    def _codes_to_wave(self, gen_ids):
        frames, i, n, resyncs = [], 0, len(gen_ids), 0
        while i <= n - 7:
            vals, ok = [], True
            for k in range(7):
                lo = self.AUDIO_OFFSET + k * 4096
                t = gen_ids[i + k]
                if not (lo <= t < lo + 4096):
                    ok = False
                    break
                vals.append(t - lo)
            if ok:
                frames.append(vals)
                i += 7
            else:
                i += 1
                resyncs += 1
        self._last_resyncs = resyncs
        if not frames:
            return None
        l1 = [f[0] for f in frames]
        l2, l3 = [], []
        for f in frames:
            l2.append(f[1]); l3.append(f[2]); l3.append(f[3])
            l2.append(f[4]); l3.append(f[5]); l3.append(f[6])
        tensors = [torch.tensor(l1)[None].to(self.device),
                   torch.tensor(l2)[None].to(self.device),
                   torch.tensor(l3)[None].to(self.device)]
        return self.snac_model.decode(tensors).squeeze().detach().cpu().numpy()

    # -------------------------------------------------------------- entry
    def __call__(self, data):
        try:
            # ---- legacy enrollment path (unchanged API) ----
            if data.get("enroll_user", False):
                pairs = data.get("enrollments", [])
                if not pairs:
                    return {"error": "No enrollment pairs provided"}
                return {"cloning_features": self.enroll_user(pairs)}

            target_text = data.get("inputs")
            if not target_text:
                return {"error": "Missing 'inputs' (target text)."}

            p = data.get("parameters", {}) or {}
            gp = {
                "temperature":        float(p.get("temperature", 0.6)),
                "top_p":              float(p.get("top_p", 0.95)),
                "top_k":              int(p.get("top_k", -1)),   # legacy default: no top-k
                "max_new_tokens":     int(p.get("max_new_tokens", 1200)),
                "repetition_penalty": float(p.get("repetition_penalty", 1.1)),
            }
            if not 0 < gp["top_p"] <= 1:
                return {"error": "top_p must be within (0, 1]."}
            if not (gp["top_k"] == -1 or gp["top_k"] > 0):
                return {"error": "top_k must be -1 (disabled) or a positive integer."}
            if not 0 < gp["repetition_penalty"] <= 2:
                return {"error": "repetition_penalty must be within (0, 2]."}
            if gp["max_new_tokens"] <= 0:
                return {"error": "max_new_tokens must be positive."}

            task   = str(p.get("task", "")).lower()
            mode   = str(p.get("mode", "vanilla")).lower()
            method = str(p.get("method", "m2")).lower()
            if mode in ("translation", "trans"):
                mode = "translate"
            if task and task not in ("tts", "vc"):
                return {"error": "parameters.task must be 'tts' or 'vc'."}
            if mode not in ("vanilla", "translate"):
                return {"error": "parameters.mode must be 'vanilla' or 'translate'."}
            if mode == "translate" and not p.get("language"):
                return {"error": "parameters.language is required for translate mode."}
            lang_tag = self._lang_display(p.get("language")) if mode == "translate" else None

            legacy_clone = bool(data.get("clone", False))
            resolved_task = "vc" if (legacy_clone or task == "vc") else "tts"

            # ---- build prompt ----
            if legacy_clone:
                # Legacy multi-pair in-context cloning (== M1 generalized)
                if data.get("clone_on_the_fly", False):
                    pairs = data.get("enrollments", [])
                    if not pairs:
                        return {"error": "No enrollment pairs provided"}
                    enrollment = []
                    for text, b64 in pairs:
                        enrollment.append({
                            "text_ids": self.tokenizer.encode(text, add_special_tokens=False),
                            "audio_codes": self._dedup_frames(
                                self._audio_to_codes(self._b64_to_wave(b64))),
                        })
                else:
                    feats = data.get("cloning_features")
                    if not feats:
                        return {"error": "No cloning features were provided"}
                    loaded = torch.load(io.BytesIO(base64.b64decode(feats)))
                    enrollment = [{
                        "text_ids": (it["text_ids"].flatten().tolist()
                                     if torch.is_tensor(it["text_ids"]) else list(it["text_ids"])),
                        "audio_codes": self._dedup_frames(list(it["audio_codes"])),
                    } for it in loaded]
                prompt_ids, method_out = [], "m1"
                for it in enrollment:
                    prompt_ids += self._text_block(it["text_ids"])
                    prompt_ids += self._audio_block(it["audio_codes"])
                prompt_ids += self._text_block(
                    self._encode_text_ids(target_text, None, lang_tag))
                prompt_ids += self._open_speech()

            elif resolved_task == "vc":
                ref_text  = p.get("reference_text")
                ref_audio = p.get("reference_audio")
                if not ref_text or not ref_audio:
                    return {"error": "VC requires parameters.reference_text and "
                                     "parameters.reference_audio (base64)."}
                if method not in ("m1", "m2"):
                    return {"error": "parameters.method must be 'm1' or 'm2'."}
                ref_codes = self._dedup_frames(self._audio_to_codes(self._b64_to_wave(ref_audio)))
                if not ref_codes:
                    return {"error": "Reference audio produced no SNAC codes."}
                tt1 = self._encode_text_ids(ref_text)
                tt2 = self._encode_text_ids(target_text, None, lang_tag)
                method_out = method
                if method == "m1":
                    prompt_ids = (self._text_block(tt1) + self._audio_block(ref_codes) +
                                  self._text_block(tt2) + self._open_speech())
                else:  # m2 continue-speaking: both texts one turn, ref codes open the AI turn
                    prompt_ids = ([self.START_OF_HUMAN, self.START_OF_TEXT] + tt1 +
                                  tt2 + [self.END_OF_TEXT, self.END_OF_HUMAN] +
                                  self._open_speech() + list(ref_codes))

            else:  # TTS (legacy default voice preserved)
                voice = p.get("voice") or p.get("speaker") or "Eniola"
                method_out = None
                prompt_ids = self._text_block(
                    self._encode_text_ids(target_text, voice, lang_tag)) + self._open_speech()

            budget = self.MAX_MODEL_LEN - gp["max_new_tokens"]
            if len(prompt_ids) > budget:
                return {"error": f"Prompt is {len(prompt_ids)} tokens; exceeds budget "
                                 f"{budget} (max_model_len - max_new_tokens)."}

            gen_ids = self._generate(prompt_ids, gp)
            wav = self._codes_to_wave(gen_ids)
            if wav is None:
                return {"error": "Model generated no audio tokens.",
                        "input_ids_len": len(prompt_ids),
                        "gen_ids_len": len(gen_ids)}

            # Legacy output format: WAV/RIFF PCM_16 base64 + raw float32 waveform
            buffer = io.BytesIO()
            sf.write(buffer, wav, samplerate=self.SNAC_SR, format="WAV", subtype="PCM_16")
            buffer.seek(0)
            audio_b64 = base64.b64encode(buffer.read()).decode("utf-8")

            return {
                "audio_sample":     wav.astype(np.float32).tolist(),   # raw waveform (legacy field)
                "audio_b64":        audio_b64,                          # base64 WAV PCM_16 (legacy)
                "sample_rate":      self.SNAC_SR,
                "input_ids_len":    len(prompt_ids),
                "gen_ids_len":      len(gen_ids),
                "duration_seconds": round(len(wav) / self.SNAC_SR, 3),
                "task": resolved_task, "mode": mode, "method": method_out,
                "decode_resyncs":   getattr(self, "_last_resyncs", 0),
                **({"gen_token_ids": gen_ids} if p.get("return_tokens") else {}),
            }

        except ValueError as e:
            return {"error": str(e)}
        except Exception as e:
            traceback.print_exc()
            return {"error": str(e)}