#!/usr/bin/env python3
"""Render the bake-off script with Chatterbox (0.5B, GPU). Chatterbox is built for
utterance-length generation, so we split into sentence groups and concatenate."""
import os, re, sys
import numpy as np
import torch, torchaudio
from chatterbox.tts import ChatterboxTTS

HERE = os.path.dirname(os.path.abspath(__file__))
script = open(os.path.join(HERE, "script.txt")).read()

# sentence-split, then greedily pack into <=280-char chunks
sentences = re.split(r"(?<=[.!?])\s+", script.replace("\n", " "))
chunks, cur = [], ""
for s in sentences:
    if len(cur) + len(s) < 280:
        cur = (cur + " " + s).strip()
    else:
        if cur: chunks.append(cur)
        cur = s
if cur: chunks.append(cur)
print(f"[chatterbox] {len(chunks)} chunks", file=sys.stderr)

model = ChatterboxTTS.from_pretrained(device="cuda")
sr = model.sr
gap = torch.zeros(1, int(sr * 0.35))
parts = []
for i, ch in enumerate(chunks):
    wav = model.generate(ch)
    parts += [wav.cpu(), gap]
    print(f"[chatterbox] {i+1}/{len(chunks)}", file=sys.stderr)

audio = torch.cat(parts, dim=1)
wav_path = os.path.join(HERE, "_cb.wav")
import soundfile as sf
sf.write(wav_path, audio.squeeze(0).numpy(), sr)

import wave, lameenc
with wave.open(wav_path) as w:
    rate, nch, frames = w.getframerate(), w.getnchannels(), w.readframes(w.getnframes())
    dur = w.getnframes() // rate
enc = lameenc.Encoder()
enc.set_bit_rate(64); enc.set_in_sample_rate(rate); enc.set_channels(nch); enc.set_quality(5)
open(os.path.join(HERE, "voice_test_chatterbox.mp3"), "wb").write(bytes(enc.encode(frames) + enc.flush()))
os.remove(wav_path)
print(f"[chatterbox] DONE {dur}s")