#!/usr/bin/env python3
"""Render podcast scripts with VibeVoice-1.5B (GPU). Runs in the vibevoice venv.

Called by daily_podcast.sh AFTER the vLLM models are stopped — VibeVoice needs the
GPU memory they hold. Loads the model once, renders every script given.

Usage: tts_vibevoice.py script1.txt [script2.txt ...]
Each scriptN.txt needs a sibling scriptN.txt.meta.json with {"title", "description",
"mp3"}; output lands in episodes/<mp3> with a matching .json for the RSS builder.
"""
import os, sys, json, wave
import torch

HERE = os.path.dirname(os.path.abspath(__file__))
EPISODES = os.path.join(HERE, "episodes")
VOICE_REF = os.path.join(HERE, "voices", "vv_voice_ref.wav")
VOICE_REF2 = os.path.join(HERE, "voices", "vv_voice_ref2.wav")
MODEL = "microsoft/VibeVoice-1.5B"

from vibevoice.modular.modeling_vibevoice_inference import (
    VibeVoiceForConditionalGenerationInference,
)
from vibevoice.processor.vibevoice_processor import VibeVoiceProcessor


import re

def clean_for_speech(s):
    """Deterministic backstop for notation the script model was told not to use.
    ASCII-only substitutions — never touches Lakota characters."""
    s = re.sub(r"\[\d+(,\s*\d+)*\]", "", s)               # citation brackets
    s = re.sub(r"https?://\S+|arXiv:\S+", "", s)           # urls / ids
    s = re.sub(r"[*_#`]+", "", s)                          # markdown leftovers
    s = re.sub(r"\$[^$]{1,80}\$", "", s)                   # inline latex
    s = re.sub(r"\b([A-Z])(\d)\b", r"\1 \2", s)            # F1 -> F 1
    s = s.replace("%", " percent")
    return re.sub(r"[ \t]{2,}", " ", s)


def render(model, processor, script, mp3_path, voice_refs=None):
    script = clean_for_speech(script)
    refs = voice_refs or [VOICE_REF]
    if len(refs) > 1:
        # dialogue script: lines already carry "Speaker 1:"/"Speaker 2:" tags.
        # Untagged continuation lines join the previous turn.
        tag = re.compile(r"^\s*speaker\s*([12])\s*:\s*(.*)$", re.I)
        turns = []
        for raw in script.split("\n"):
            if not raw.strip():
                continue
            m = tag.match(raw)
            if m:
                turns.append([int(m.group(1)), m.group(2).strip()])
            elif turns:
                turns[-1][1] += " " + raw.strip()
            else:
                turns.append([1, raw.strip()])
        text = "\n".join(f"Speaker {n}: {t}" for n, t in turns)
    else:
        # one "Speaker 1: ..." LINE per paragraph — the processor only parses the
        # tag at line start; a mid-line tag is read aloud as the words "Speaker one"
        paras = [p.strip().replace("\n", " ") for p in script.split("\n\n") if p.strip()]
        text = "\n".join("Speaker 1: " + p for p in paras)
    inputs = processor(text=[text], voice_samples=[refs],
                       padding=True, return_tensors="pt", return_attention_mask=True)
    inputs = {k: (v.to("cuda") if hasattr(v, "to") else v) for k, v in inputs.items()}
    with torch.no_grad():
        out = model.generate(**inputs, max_new_tokens=None, cfg_scale=1.3,
                             tokenizer=processor.tokenizer,
                             generation_config={"do_sample": False}, verbose=False)
    wav_path = mp3_path.replace(".mp3", ".wav")
    processor.save_audio(out.speech_outputs[0], output_path=wav_path)
    import lameenc
    with wave.open(wav_path) as w:
        rate, nch = w.getframerate(), w.getnchannels()
        frames, dur = w.readframes(w.getnframes()), w.getnframes() // w.getframerate()
    enc = lameenc.Encoder()
    enc.set_bit_rate(64); enc.set_in_sample_rate(rate); enc.set_channels(nch); enc.set_quality(5)
    open(mp3_path, "wb").write(bytes(enc.encode(frames) + enc.flush()))
    os.remove(wav_path)
    return dur


def main():
    jobs = [p for p in sys.argv[1:] if os.path.exists(p) and os.path.exists(p + ".meta.json")]
    if not jobs:
        print("[vv-tts] no scripts to render")
        return
    print(f"[vv-tts] loading {MODEL}...", file=sys.stderr)
    processor = VibeVoiceProcessor.from_pretrained(MODEL)
    model = VibeVoiceForConditionalGenerationInference.from_pretrained(
        MODEL, torch_dtype=torch.bfloat16, device_map="cuda", attn_implementation="sdpa")
    model.eval()
    for p in jobs:
        meta = json.load(open(p + ".meta.json"))
        mp3 = os.path.join(EPISODES, meta["mp3"])
        # render to a temp name, promote only when complete — the RSS builder must
        # never see a half-written MP3 (it skips any '.part.' file)
        tmp_mp3 = mp3.replace(".mp3", ".part.mp3")
        refs = [VOICE_REF, VOICE_REF2][:meta.get("speakers", 1)]
        dur = render(model, processor, open(p).read(), tmp_mp3, voice_refs=refs)
        os.replace(tmp_mp3, mp3)
        json.dump({"title": meta["title"], "description": meta["description"]},
                  open(mp3.replace(".mp3", ".json"), "w"))
        if "paper_id" in meta:   # success: NOW the paper counts as featured
            fp = meta["featured_path"]
            feat = set(json.load(open(fp))) if os.path.exists(fp) else set()
            feat.add(meta["paper_id"])
            json.dump(sorted(feat), open(fp, "w"))
        print(f"[vv-tts] {meta['mp3']} ({dur}s)")


if __name__ == "__main__":
    main()