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

Task matrix (routed by `parameters`):
  task="tts", mode="vanilla"    : {speaker}: text                      -> speech
  task="tts", mode="translate"  : {speaker} - {Language}: text         -> speech in Language
  task="vc",  mode="vanilla"    : reference (text+audio) + target text -> speech in reference voice
  task="vc",  mode="translate"  : + language tag on target text        -> cross-lingual cloning
  VC method="m1" (in-context) | method="m2" (continue-speaking)

Output parity with the legacy endpoint: `audio_b64` is base64 of the RAW
float32 little-endian mono PCM buffer at 24000 Hz (NO WAV/RIFF container),
so existing products decode with:  np.frombuffer(base64.b64decode(s), dtype=np.float32)
[If the legacy endpoint used int16, change RAW_DTYPE to np.int16 below.]

Prompts are byte-identical to Step-III training (_encode_text / build_tts /
build_vc_both), reference codes are frame-deduped, and prompts reach vLLM as
token ids (never a decoded string).
"""

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:
    # ---- Orpheus special tokens (fixed by the model) ----
    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

    # NOTE: fine-tune data capped at 2048 tokens; 4096 kept so M1-VC prompts
    # (ref codes + two texts, often 1000-2000 tokens) retain a generation
    # budget. Base Llama-3 RoPE supports these positions natively; expect the
    # best quality when prompt+generation stays near the trained ~2048.
    MAX_MODEL_LEN   = 4096
    MAX_REF_SECONDS = 30
    SNAC_SR         = 24000
    RAW_DTYPE       = np.float32          # legacy raw-PCM dtype (see docstring)

    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"
        # SNAC first (tiny, ~80 MB) so it never contends with vLLM's reservation.
        self.snac = 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,
        )
        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(self, text, speaker=None, lang_tag=None, add_bos=True):
        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
        ids = self.tokenizer.encode(prompt, add_special_tokens=add_bos)
        ids.append(self.END_OF_TEXT)
        return ids

    # ------------------------------------------------------ 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:
            # temp-file fallback: librosa/audioread handles containers
            # libsndfile can't open, but needs a real file path for some codecs.
            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                    # librosa returns (ch, n)
            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. "
                             f"Send a shorter clip.")
        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.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 = 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

    # ------------------------------------------------------ prompt builders
    def build_tts_prompt(self, text, speaker, mode, language):
        lang_tag = self._lang_display(language) if mode == "translate" else None
        tt = self._encode_text(text, speaker, lang_tag, add_bos=True)
        return [self.START_OF_HUMAN] + tt + [self.END_OF_HUMAN,
                self.START_OF_AI, self.START_OF_SPEECH]

    def build_vc_prompt(self, ref_text, ref_codes, target_text, mode, language, method):
        tag2 = self._lang_display(language) if mode == "translate" else None
        tt1 = self._encode_text(ref_text, None, None, add_bos=True)
        tt2 = self._encode_text(target_text, None, tag2, add_bos=False)
        if method == "m1":
            return ([self.START_OF_HUMAN] + tt1 + [self.END_OF_HUMAN,
                     self.START_OF_AI, self.START_OF_SPEECH] + ref_codes +
                    [self.END_OF_SPEECH, self.END_OF_AI,
                     self.START_OF_HUMAN] + tt2 + [self.END_OF_HUMAN,
                     self.START_OF_AI, self.START_OF_SPEECH])
        return ([self.START_OF_HUMAN] + tt1 + tt2 + [self.END_OF_HUMAN,
                 self.START_OF_AI, self.START_OF_SPEECH] + ref_codes)

    # --------------------------------------------------------- generation
    def _generate(self, prompt_ids, params):
        sampling = SamplingParams(
            temperature        = params["temperature"],
            top_p              = params["top_p"],
            top_k              = params["top_k"],
            max_tokens         = params["max_new_tokens"],
            repetition_penalty = params["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):
        """Frame-validating SNAC decode: accepts only well-formed 7-token frames
        (token k in slot-k range), resyncs on malformed spans."""
        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, 0
        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),
        ]
        wav = self.snac.decode(tensors).squeeze().detach().cpu().numpy()
        return wav, len(frames)

    def _wave_to_b64_raw(self, wav):
        """Legacy parity: base64 of raw little-endian PCM buffer, no container."""
        return base64.b64encode(
            np.ascontiguousarray(wav.astype(self.RAW_DTYPE)).tobytes()
        ).decode("utf-8")

    # -------------------------------------------------------------- entry
    def __call__(self, data):
        try:
            target_text = data.get("inputs")
            if not target_text:
                return {"error": "Missing 'inputs' (target text)."}

            p = data.get("parameters", {}) or {}
            task   = str(p.get("task", "tts")).lower()
            mode   = str(p.get("mode", "vanilla")).lower()
            method = str(p.get("method", "m2")).lower()
            if mode in ("translation", "trans"):
                mode = "translate"

            if 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."}

            gen_params = {
                "temperature":        float(p.get("temperature", 0.6)),
                "top_p":              float(p.get("top_p", 0.95)),
                "top_k":              int(p.get("top_k", 50)),
                "max_new_tokens":     int(p.get("max_new_tokens", 1200)),
                "repetition_penalty": float(p.get("repetition_penalty", 1.1)),
            }
            if not 0 < gen_params["top_p"] <= 1:
                return {"error": "top_p must be within (0, 1]."}
            if not (gen_params["top_k"] == -1 or gen_params["top_k"] > 0):
                return {"error": "top_k must be -1 (disabled) or a positive integer."}
            if not 0 < gen_params["repetition_penalty"] <= 2:
                return {"error": "repetition_penalty must be within (0, 2]."}
            if gen_params["max_new_tokens"] <= 0:
                return {"error": "max_new_tokens must be positive."}

            if 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_wave  = self._b64_to_wave(ref_audio)
                ref_codes = self._dedup_frames(self._audio_to_codes(ref_wave))
                if not ref_codes:
                    return {"error": "Reference audio produced no SNAC codes."}
                prompt_ids = self.build_vc_prompt(
                    ref_text, ref_codes, target_text, mode, p.get("language"), method)
            else:
                prompt_ids = self.build_tts_prompt(
                    target_text, p.get("voice") or p.get("speaker"),
                    mode, p.get("language"))

            budget = self.MAX_MODEL_LEN - gen_params["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). "
                                 f"Shorten the reference clip or text."}

            gen_ids = self._generate(prompt_ids, gen_params)
            wav, n_frames = self._codes_to_wave(gen_ids)
            if wav is None:
                return {"error": "Model generated no audio tokens.",
                        "input_tokens": len(prompt_ids),
                        "generated_tokens": len(gen_ids)}

            return {
                "audio_b64":        self._wave_to_b64_raw(wav),   # RAW float32 PCM (legacy parity)
                "audio_dtype":      np.dtype(self.RAW_DTYPE).name,
                "sample_rate":      self.SNAC_SR,
                "duration_seconds": round(len(wav) / self.SNAC_SR, 3),
                "audio_frames":     n_frames,
                "input_tokens":     len(prompt_ids),
                "generated_tokens": len(gen_ids),
                "task": task, "mode": mode,
                "method": method if task == "vc" else None,
                "decode_resyncs":   getattr(self, "_last_resyncs", 0),
            }

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