# worker_tts.py — stage 8 TTS. FastAPI :8006. dots.tts loads in-process (session-guarded).
# generates audio locally, uploads to minio, returns uri + duration. no local persistence.
import os, re, json, uuid, wave, contextlib, subprocess, logging, functools
from fastapi import FastAPI
from pydantic import BaseModel
import transport

log = logging.getLogger("tts")

app = FastAPI()
transport.install_logging(app, "tts")
SHM = "/dev/shm"
DOTS_MODEL = "rednote-hilab/dots.tts-base"
_tts = None

# B#3 one fixed voice: dots.tts samples a RANDOM speaker each call unless given a reference clip,
# so the narrator's timbre drifts panel-to-panel. clone from ONE reference for every synth.
# to pin YOUR OWN voice: set VOICE_REF=/path/to/clip.wav + VOICE_REF_TEXT="its exact transcript"
#   (5-15s of clean speech works best). no edit/rebuild needed -- just the env vars on the worker.
# with no VOICE_REF, we bootstrap a seeded reference once and persist it; delete it to reroll.
VOICE_DIR = os.path.expanduser("~/.cache/manga-tts")
REF_WAV = os.environ.get("VOICE_REF") or os.path.join(VOICE_DIR, "narrator_ref.wav")
REF_TEXT = os.environ.get("VOICE_REF_TEXT", "The story continues as our hero steps forward into the unknown.")
REF_SEED = 20260713


def _load_tts():
    global _tts
    if _tts is None:
        # dots.tts's vendored loader (models/dots_tts/model.py) calls AutoTokenizer.from_pretrained
        # with no kwargs, so its Mistral-derived tokenizer loads with the buggy split regex. Default
        # fix_mistral_regex=True at the transformers layer to get canonical tokenization + kill the warning.
        # ponytail: monkeypatch because the loader exposes no passthrough; drop if dots_tts adds one.
        import transformers
        _orig = transformers.AutoTokenizer.from_pretrained.__func__
        transformers.AutoTokenizer.from_pretrained = classmethod(
            lambda cls, *a, **kw: _orig(cls, *a, **{"fix_mistral_regex": True, **kw})
        )
        # dots_tts/__init__.py refuses to import when torch and torchaudio minors differ. workpc runs
        # Arch's torch 2.13 but pytorch.org ships no torchaudio past 2.11 for ROCm, so the pair can't
        # be satisfied; 2.11 loads and resamples fine against 2.13. Lie to the guard for the import.
        # ponytail: drop this once a torchaudio matching torch's minor exists for ROCm.
        import torch, importlib.metadata as _md
        _ver = _md.version
        _md.version = lambda n: torch.__version__ if n == "torchaudio" else _ver(n)
        try:
            from dots_tts.runtime import DotsTtsRuntime
        finally:
            _md.version = _ver
        _tts = DotsTtsRuntime.from_pretrained(DOTS_MODEL, precision="bfloat16")
    return _tts


def _wav_duration(path: str) -> float:
    with contextlib.closing(wave.open(path, "rb")) as w:
        return round(w.getnframes() / float(w.getframerate()), 3)


class TTSInput(BaseModel):
    text: str
    speaker: str = "narrator"   # multi-voice is v3
    panel_id: str = ""
    session_id: str = ""
    panel_uri: str = ""         # optional: if passed, audio is stored beside its panel


def _audio_uri(data: "TTSInput") -> str:
    # prefer the panel's own manga/chapter prefix; the orchestrator currently doesn't pass it,
    # so fall back to a flat panel_id-keyed key (matches homesrv's panel_id-keyed audio table).
    # ponytail: flat key collides across chapters (as does the homesrv audio table); pass
    #   panel_uri from run_stage_tts to make it per-chapter unique.
    if data.panel_uri:
        manga_id, chapter_id = transport.ids_from_uri(data.panel_uri)
        return transport.AUDIO_URI.format(manga_id=manga_id, chapter_id=chapter_id,
                                          name=data.panel_id or 'p')
    return transport.AUDIO_FLAT_URI.format(name=data.panel_id or 'p')


def _ensure_ref() -> str:
    """the fixed narrator reference clip; generate it once (seeded) and persist across restarts."""
    if os.path.exists(REF_WAV):
        return REF_WAV
    os.makedirs(VOICE_DIR, exist_ok=True)
    try:
        import torch
        torch.manual_seed(REF_SEED)     # reproducible speaker for the bootstrap sample
    except Exception:
        pass
    out = _load_tts().generate(text=REF_TEXT)
    _write_wav(out["audio"], out["sample_rate"], REF_WAV)
    return REF_WAV


def _calm(text: str) -> str:
    """dots.tts over-emotes on '!' (shouty prosody). soften exclamations to periods so the narrator
    stays even. only the spoken text is calmed -- the burned subtitles keep the original '!'."""
    return re.sub(r"\s*!+", ".", text)


# 158: dots.tts has no SSML/phoneme input, so proper nouns it mangles are fixed by respelling the SPOKEN
# text only (burned subtitles keep the original spelling — they're built elsewhere from the untouched
# script). Lazy v1: one global JSON map {term: phonetic}, ~8 lines, loaded once.
LEXICON_PATH = os.path.expanduser(os.environ.get("TTS_LEXICON", "~/.cache/manga-tts/lexicon.json"))


@functools.lru_cache(maxsize=1)
def _lexicon():
    """(compiled word-boundary pattern, {lower_term: phonetic}) or None. Cached; delete the file and
    call _lexicon.cache_clear() to reload. Longest terms first so multi-word names match whole."""
    try:
        with open(LEXICON_PATH) as f:
            m = json.load(f)
    except (OSError, ValueError):
        return None
    m = {k: v for k, v in (m or {}).items() if k and v}
    if not m:
        return None
    pat = re.compile(r"\b(" + "|".join(re.escape(k) for k in sorted(m, key=len, reverse=True)) + r")\b",
                     re.IGNORECASE)
    return pat, {k.lower(): v for k, v in m.items()}


def _respell(text: str) -> str:
    """Substitute known proper nouns with their phonetic respelling (whole word, case-insensitive)."""
    lex = _lexicon()
    if not lex:
        return text
    pat, lookup = lex
    return pat.sub(lambda mo: lookup[mo.group(0).lower()], text)


def _generate(text: str) -> str:
    """returns a local wav path. dots runtime returns {"audio": samples, "sample_rate": sr};
    clone the fixed reference voice so every panel narrates in the same timbre."""
    ref = _ensure_ref()
    out = _load_tts().generate(text=_respell(_calm(text)), prompt_audio_path=ref, prompt_text=REF_TEXT)
    return _write_wav(out["audio"], out["sample_rate"])


def _write_wav(audio, sample_rate: int, path: str | None = None) -> str:
    import numpy as np, soundfile as sf
    if hasattr(audio, "detach"):        # torch tensor (possibly on GPU)
        audio = audio.detach().cpu().numpy()
    a = np.asarray(audio, dtype="float32").squeeze()   # (samples,) mono
    if path is None:
        path = f"{SHM}/tts_{uuid.uuid4().hex[:8]}.wav"
    sf.write(path, a, sample_rate, subtype="PCM_16")
    return path


def _loudnorm(path: str) -> str:
    """EBU R128 loudness-normalize so narration volume is even panel-to-panel (#14).
    Returns a normalized path; on any ffmpeg failure returns the original (never lose audio)."""
    sr = wave.open(path, "rb").getframerate()
    out = f"{SHM}/ln_{uuid.uuid4().hex[:8]}.wav"
    try:
        subprocess.run(
            ["ffmpeg", "-y", "-i", path, "-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
             "-ar", str(sr), out],
            check=True, capture_output=True,
        )
        os.replace(out, path)
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        log.warning("loudnorm skipped for %s: %r", path, e)
        if os.path.exists(out):
            os.remove(out)
    return path


@app.post("/tts")
async def tts(data: TTSInput):
    local = _loudnorm(_generate(data.text))
    uri = _audio_uri(data)
    transport.put(local, uri)
    dur = _wav_duration(local)
    os.remove(local)
    return {"audio_uri": uri, "duration": dur}


@app.post("/unload")
async def unload():
    """free the resident dots.tts so the session manager can hand the GPU to the next model."""
    global _tts
    was = _tts is not None
    _tts = None
    import gc; gc.collect()
    try:
        import torch; torch.cuda.empty_cache()
    except Exception:
        pass
    return {"ok": True, "unloaded": was}


@app.get("/health")
async def health():
    return {"status": "ok"}


if __name__ == "__main__":
    # self-check: _write_wav renders a sample array (1.0s @16k) to a PCM-16 wav that
    # wave.open reads back at the right duration; then exercise upload + uri.
    import math

    class _FakeMC:
        def __init__(self): self.store = {}
        def bucket_exists(self, b): return True
        def make_bucket(self, b): pass
        def fput_object(self, b, k, path): self.store[(b, k)] = open(path, "rb").read()

    samples = [0.3 * math.sin(i / 8) for i in range(16000)]   # bare python list -> np.asarray
    local = _write_wav(samples, 16000)
    dur = _wav_duration(local)
    uri = transport.put(local, "s3://manga/m/c/audio/p001.wav", client=_FakeMC())
    assert uri.endswith("audio/p001.wav") and abs(dur - 1.0) < 0.01, (uri, dur)
    os.remove(local)
    explicit = f"{SHM}/tts_selfcheck_ref.wav"          # _write_wav honors an explicit path (ref clip)
    assert _write_wav(samples, 16000, explicit) == explicit and os.path.exists(explicit)
    os.remove(explicit)
    assert _calm("Stop!! Now!") == "Stop. Now." and _calm("no bangs") == "no bangs"  # even prosody

    # _audio_uri: a panel uri puts the wav beside its chapter, no panel uri falls back to the flat key.
    assert _audio_uri(TTSInput(text="x", panel_id="p003",
                               panel_uri="s3://panels/m1/c1/panels/p003.png")) \
        == "s3://audio/m1/c1/audio/p003.wav"
    assert _audio_uri(TTSInput(text="x", panel_id="p003")) == "s3://audio/_audio/p003.wav"
    assert _audio_uri(TTSInput(text="x")) == "s3://audio/_audio/p.wav"

    # 158: pronunciation lexicon respells whole words only, case-insensitive, spoken text only.
    import tempfile, json as _json
    globals()["LEXICON_PATH"] = os.path.join(tempfile.mkdtemp(), "lex.json")
    with open(LEXICON_PATH, "w") as f:
        _json.dump({"Najimi": "nah-jee-mee", "Choi Haeseon": "chwe hae-son"}, f)
    _lexicon.cache_clear()
    assert _respell("Then Najimi ran.") == "Then nah-jee-mee ran."          # single name
    assert _respell("with najimi today") == "with nah-jee-mee today"        # case-insensitive
    assert _respell("Najimist stays") == "Najimist stays"                   # word boundary (no substring)
    assert _respell("Choi Haeseon smiled") == "chwe hae-son smiled"         # multi-word term
    os.remove(LEXICON_PATH); _lexicon.cache_clear()
    assert _respell("Najimi ran.") == "Najimi ran."                          # no file -> passthrough

    # loudnorm: with ffmpeg present the wav is normalized in place and stays readable at its sr;
    # without ffmpeg it's a safe no-op returning the same path (audio never lost).
    ln = _write_wav(samples, 16000)
    have_ffmpeg = subprocess.run(["ffmpeg", "-version"], capture_output=True, check=False).returncode == 0 \
        if __import__("shutil").which("ffmpeg") else False
    assert _loudnorm(ln) == ln and os.path.exists(ln)
    assert wave.open(ln, "rb").getframerate() == 16000
    os.remove(ln)
    print("worker_tts self-check ok" + ("" if have_ffmpeg else " (ffmpeg absent, loudnorm no-op tested)"))