#!/usr/bin/env python3
"""mOrpheus TTS wrapper — turns the vLLM engine's Orpheus audio tokens into 24kHz WAV.

Architecture: this CPU service builds the Orpheus prompt (as raw token ids), calls the
vLLM engine (which serves the mOrpheus LLM), recovers the generated audio-token ids by
re-tokenizing the returned text (vLLM emits them as `<custom_token_N>` strings when
skip_special_tokens=false), then SNAC-decodes them to audio.

Token scheme (verified): audio-base 128266, 7-token SNAC frames (pos0->L1, pos1/4->L2,
pos2/3/5/6->L3, offsets k*4096). Named-speaker prompt: [SOH] "voice: text" [EOT][SOA].
Zero-shot cloning: [BOS][SOH] ref_text [EOT][SOA][SOS] <ref audio tokens> [EOS_sp] then
[SOH] text [EOT][SOA]. Generation stops at end-of-speech (128258).
"""
import os, io, re, json, base64
import numpy as np, torch, soundfile as sf, requests
from scipy.signal import resample_poly
from fastapi import FastAPI, HTTPException
from fastapi.responses import Response, StreamingResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from transformers import AutoTokenizer
from snac import SNAC

MODEL_DIR = os.environ.get("MORPHEUS_MODEL_DIR", "/model")
SNAC_DIR = os.environ.get("SNAC_DIR", "/snac")
VLLM_URL = os.environ.get("VLLM_URL", "http://vllm-morpheus:8000/v1/completions")
VLLM_MODEL = os.environ.get("VLLM_MODEL", "morpheus")
SNAC_DEVICE = os.environ.get("SNAC_DEVICE", "cpu")
DEFAULT_VOICE = os.environ.get("MORPHEUS_DEFAULT_VOICE", "baddy")
VOICES = [v for v in os.environ.get("MORPHEUS_VOICES", "baddy").split(",") if v]
AUDIO_MAX = 156937
AUDIO_BASE, SOS, EOS_SP, SOH, SOA, EOT, BOS = 128266, 128257, 128258, 128259, 128260, 128009, 128000
MAX_CTX = int(os.environ.get("MORPHEUS_MAX_LEN", "4096"))  # must match the engine --max-model-len


def _cap(prompt_ids: list[int], want: int) -> int:
    # never let prompt + gen exceed the context (a cloning ref block is ~1100 tokens)
    return max(64, min(want, MAX_CTX - len(prompt_ids) - 16))

tok = AutoTokenizer.from_pretrained(MODEL_DIR)
snac_model = SNAC.from_pretrained(SNAC_DIR).to(SNAC_DEVICE).eval()
app = FastAPI(title="mOrpheus TTS", version="0.1.0")
# INTERNAL RESEARCH: browser callers (e.g. gateway-chat on ana-docker:8091) fetch this
# cross-origin; allow all origins on this internal-only endpoint.
app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_methods=["*"],
                   allow_headers=["*"], expose_headers=["X-Audio-Seconds"])


class TTSReq(BaseModel):
    text: str
    voice: str = DEFAULT_VOICE
    temperature: float = 0.6
    top_p: float = 0.95
    max_tokens: int = 3500           # ~42s of audio (auto-clamped to fit MAX_CTX; see _cap)
    repetition_penalty: float = 1.1  # LOAD-BEARING: rep 1.0 => model never stops (rambles to
                                     # the cap); rep 1.1 => clean end-of-speech. Keep <=1.1 for
                                     # cloning (higher penalizes the in-context ref audio tokens).
    reference_audio_b64: str | None = None  # optional zero-shot clone: base64 WAV
    reference_text: str | None = None       # transcript of the reference


def _encode_ref(wav_bytes: bytes, ref_text: str) -> list[int]:
    wav, sr = sf.read(io.BytesIO(wav_bytes))
    if wav.ndim > 1:
        wav = wav.mean(1)
    wav = wav.astype(np.float32)
    if sr != 24000:
        wav = resample_poly(wav, 24000, sr).astype(np.float32)
    wt = torch.tensor(wav, device=SNAC_DEVICE).view(1, 1, -1)
    with torch.inference_mode():
        codes = snac_model.encode(wt)
    L1, L2, L3 = [c.squeeze(0).tolist() for c in codes]
    ids = []
    for i in range(len(L1)):
        ids += [L1[i] + AUDIO_BASE, L2[2 * i] + AUDIO_BASE + 4096, L3[4 * i] + AUDIO_BASE + 8192,
                L3[4 * i + 1] + AUDIO_BASE + 12288, L2[2 * i + 1] + AUDIO_BASE + 16384,
                L3[4 * i + 2] + AUDIO_BASE + 20480, L3[4 * i + 3] + AUDIO_BASE + 24576]
    return [BOS, SOH] + tok(ref_text, add_special_tokens=False).input_ids + [EOT, SOA, SOS] + ids + [EOS_SP]


# Pre-staged clone voices: a /voices dir of <name>.wav + <name>.txt (its transcript). Each is
# encoded once to its Orpheus reference block at startup, so voice="<name>" zero-shot clones it.
VOICES_DIR = os.environ.get("MORPHEUS_VOICES_DIR", "/voices")
CLONE_REFS = {}
if os.path.isdir(VOICES_DIR):
    for fn in sorted(os.listdir(VOICES_DIR)):
        if fn.endswith(".wav") and os.path.exists(os.path.join(VOICES_DIR, fn[:-4] + ".txt")):
            name = fn[:-4]
            try:
                CLONE_REFS[name] = _encode_ref(open(os.path.join(VOICES_DIR, fn), "rb").read(),
                                               open(os.path.join(VOICES_DIR, name + ".txt")).read().strip())
                print(f"[voices] staged clone voice '{name}'", flush=True)
            except Exception as e:
                print(f"[voices] FAILED to stage '{name}': {e}", flush=True)


def _build_prompt(req: TTSReq) -> list[int]:
    if req.reference_audio_b64 and req.reference_text:
        ref = _encode_ref(base64.b64decode(req.reference_audio_b64), req.reference_text)
        return ref + [SOH] + tok(req.text, add_special_tokens=False).input_ids + [EOT, SOA]
    if req.voice in CLONE_REFS:                      # pre-staged clone voice (e.g. beatrice)
        return CLONE_REFS[req.voice] + [SOH] + tok(req.text, add_special_tokens=False).input_ids + [EOT, SOA]
    return [SOH] + tok(f"{req.voice}: {req.text}").input_ids + [EOT, SOA]


def _decode(ids: list[int]):
    if SOS in ids:
        ids = ids[len(ids) - 1 - ids[::-1].index(SOS) + 1:]
    codes = [t - AUDIO_BASE for t in ids if AUDIO_BASE <= t <= AUDIO_MAX]
    l1, l2, l3 = [], [], []
    for i in range(len(codes) // 7):
        f = codes[7 * i:7 * i + 7]
        c = [f[0], f[1] - 4096, f[2] - 8192, f[3] - 12288, f[4] - 16384, f[5] - 20480, f[6] - 24576]
        if any(x < 0 or x > 4095 for x in c):
            continue
        l1.append(c[0]); l2 += [c[1], c[4]]; l3 += [c[2], c[3], c[5], c[6]]
    if not l1:
        return None
    ct = [torch.tensor(l1).unsqueeze(0).to(SNAC_DEVICE),
          torch.tensor(l2).unsqueeze(0).to(SNAC_DEVICE),
          torch.tensor(l3).unsqueeze(0).to(SNAC_DEVICE)]
    with torch.inference_mode():
        return snac_model.decode(ct).squeeze().cpu().numpy()


@app.get("/health")
def health():
    return {"status": "ok", "voices": VOICES + list(CLONE_REFS), "default": DEFAULT_VOICE, "engine": VLLM_URL}


@app.get("/voices")
def voices():
    return {"voices": VOICES + list(CLONE_REFS), "cloned": list(CLONE_REFS), "default": DEFAULT_VOICE}


@app.post("/tts")
def tts(req: TTSReq):
    prompt_ids = _build_prompt(req)
    payload = {"model": VLLM_MODEL, "prompt": prompt_ids, "max_tokens": _cap(prompt_ids, req.max_tokens),
               "temperature": req.temperature, "top_p": req.top_p, "skip_special_tokens": False,
               "stop_token_ids": [EOS_SP], "repetition_penalty": req.repetition_penalty}
    try:
        r = requests.post(VLLM_URL, json=payload, timeout=180)
    except requests.RequestException as e:
        raise HTTPException(502, f"engine unreachable: {e}")
    if r.status_code != 200:
        raise HTTPException(502, f"engine {r.status_code}: {r.text[:200]}")
    text = r.json()["choices"][0]["text"]
    gen_ids = tok(text, add_special_tokens=False).input_ids
    audio = _decode(gen_ids)
    if audio is None:
        raise HTTPException(500, "no audio tokens generated")
    buf = io.BytesIO()
    sf.write(buf, audio, 24000, format="WAV", subtype="PCM_16")
    return Response(content=buf.getvalue(), media_type="audio/wav",
                    headers={"X-Audio-Seconds": f"{len(audio)/24000:.2f}"})


# --- streaming decode: emit raw PCM16 (24kHz mono) as the engine generates -----------
# vLLM emits audio tokens as <custom_token_N> strings (token_id = 128256 + N). We decode in
# WINDOWED CHUNKS: every CHUNK new frames, decode a small window [CTX left | CHUNK | CTX
# right] and emit only the middle CHUNK frames (context on both sides => seamless, no
# clicks). O(window) per call and few calls => decode keeps up with gen on CPU (per-frame
# decode did NOT — per-call overhead x ~60 frames serialized to seconds). TTFA ~0.7s.
CUSTOM_RE = re.compile(r"<custom_token_(\d+)>")
FRAME = 2048  # audio samples per SNAC frame


def _decode_frames(ids: list[int]):
    codes = [t - AUDIO_BASE for t in ids if AUDIO_BASE <= t <= AUDIO_MAX]
    n = len(codes) // 7
    if n == 0:
        return None
    l1, l2, l3 = [], [], []
    for i in range(n):
        f = codes[7 * i:7 * i + 7]
        c = [f[0], f[1] - 4096, f[2] - 8192, f[3] - 12288, f[4] - 16384, f[5] - 20480, f[6] - 24576]
        if any(x < 0 or x > 4095 for x in c):
            return None
        l1.append(c[0]); l2 += [c[1], c[4]]; l3 += [c[2], c[3], c[5], c[6]]
    ct = [torch.tensor(l1).unsqueeze(0).to(SNAC_DEVICE),
          torch.tensor(l2).unsqueeze(0).to(SNAC_DEVICE),
          torch.tensor(l3).unsqueeze(0).to(SNAC_DEVICE)]
    with torch.inference_mode():
        return snac_model.decode(ct).squeeze().cpu().numpy()


def _pcm16(arr) -> bytes:
    return (np.clip(arr, -1, 1) * 32767).astype(np.int16).tobytes()


@app.post("/tts/stream")
def tts_stream(req: TTSReq):
    prompt_ids = _build_prompt(req)
    payload = {"model": VLLM_MODEL, "prompt": prompt_ids, "max_tokens": _cap(prompt_ids, req.max_tokens),
               "temperature": req.temperature, "top_p": req.top_p, "skip_special_tokens": False,
               "stop_token_ids": [EOS_SP], "repetition_penalty": req.repetition_penalty, "stream": True}
    CTX, CHUNK = 2, 6  # context frames each side, frames emitted per decode

    def emit(emitted, total, final):
        # decode [emitted-CTX : end] with left/right context; return (pcm_bytes, new_emitted)
        end = total if final else emitted + CHUNK
        ws = max(0, emitted - CTX)
        audio = _decode_frames(buf[7 * ws:7 * (end + (0 if final else CTX))])
        if audio is None:
            return b"", end
        off = (emitted - ws) * FRAME
        chunk = audio[off:] if final else audio[off:off + CHUNK * FRAME]
        return (_pcm16(chunk) if len(chunk) else b""), end

    buf = []

    def gen():
        acc, cursor, emitted = "", 0, 0
        with requests.post(VLLM_URL, json=payload, stream=True, timeout=180) as r:
            for raw in r.iter_lines():
                if not raw:
                    continue
                s = raw.decode("utf-8", "ignore")
                if not s.startswith("data:"):
                    continue
                d = s[5:].strip()
                if d == "[DONE]":
                    break
                try:
                    j = json.loads(d)
                except Exception:
                    continue
                acc += j["choices"][0].get("text", "")
                for m in CUSTOM_RE.finditer(acc, cursor):
                    cursor = m.end()
                    tid = 128256 + int(m.group(1))
                    if AUDIO_BASE <= tid <= AUDIO_MAX:
                        buf.append(tid)
                while (len(buf) // 7) - emitted >= CHUNK + CTX:   # enough new frames + right ctx
                    b, emitted = emit(emitted, len(buf) // 7, final=False)
                    if b:
                        yield b
        if (len(buf) // 7) > emitted:                             # flush the tail
            b, emitted = emit(emitted, len(buf) // 7, final=True)
            if b:
                yield b

    return StreamingResponse(gen(), media_type="application/octet-stream",
                             headers={"Cache-Control": "no-cache", "X-Sample-Rate": "24000"})
