#!/usr/bin/env python3
"""
speak.py
Generate speech in Brian's voice from text, using Chatterbox (MIT-licensed, runs locally).

Long scripts are split into sentence-sized pieces, each piece is generated in the voice of
the reference clip, and the pieces are joined with natural pauses into one WAV file.
A settings file is saved next to every output so any take can be reproduced exactly.

Examples:
  python speak.py "Hey guys, welcome back to the Reptarium." --ref brian_voice_refs\\ref_01.wav
  python speak.py --file script.txt --ref brian_voice_refs\\ref_01.wav --takes 3
  python speak.py --interactive --ref brian_voice_refs\\ref_01.wav
  python speak.py "Hey guys." --voice voices\\brian.pt     (profile built from all clips)

Every output carries Resemble AI's inaudible Perth watermark, which marks it as AI-generated.
"""

import os
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")  # Windows: avoid symlink errors
os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import argparse
import json
import random
import re
import sys
import time
import warnings
from datetime import datetime
from pathlib import Path

warnings.filterwarnings("ignore")  # the libraries print a lot of harmless deprecation noise

import numpy as np
import soundfile as sf


# ---------------------------------------------------------------- text splitting

def split_text(text, max_chars=280):
    """Split a script into chunks Chatterbox handles well.

    Returns a list of (chunk, pause_after) where pause_after is "sentence" or "paragraph".
    Chatterbox gets unreliable on very long inputs, so long sentences are split at commas,
    then at word boundaries as a last resort."""
    paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text.strip()) if p.strip()]
    out = []
    for para in paragraphs:
        para = re.sub(r"\s+", " ", para)
        sentences = [s.strip() for s in re.split(r"(?<=[.!?])\s+", para) if s.strip()]
        pieces = []
        for s in sentences:
            if len(s) <= max_chars:
                pieces.append(s)
                continue
            cur = ""
            for part in re.split(r"(?<=[,;:])\s+", s):
                while len(part) > max_chars:  # no commas to break on: cut at a word boundary
                    cut = part.rfind(" ", 0, max_chars)
                    cut = cut if cut > 0 else max_chars
                    if cur:
                        pieces.append(cur)
                        cur = ""
                    pieces.append(part[:cut].strip())
                    part = part[cut:].strip()
                if cur and len(cur) + 1 + len(part) > max_chars:
                    pieces.append(cur)
                    cur = part
                else:
                    cur = f"{cur} {part}".strip()
            if cur:
                pieces.append(cur)
        # Join very short sentences ("Look at this. Wow.") so they flow naturally.
        merged = []
        for piece in pieces:
            if merged and len(merged[-1]) < 60 and len(merged[-1]) + 1 + len(piece) <= max_chars:
                merged[-1] = f"{merged[-1]} {piece}"
            else:
                merged.append(piece)
        for i, m in enumerate(merged):
            out.append((m, "paragraph" if i == len(merged) - 1 else "sentence"))
    return out


# ---------------------------------------------------------------- model

def pick_device(choice):
    import torch
    if choice != "auto":
        return choice
    if torch.cuda.is_available():
        return "cuda"
    if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def load_model(kind, device):
    print(f"Loading Chatterbox ({kind}) on {device}. First run downloads the model (~2-3 GB)...")
    t = time.time()
    if kind == "turbo":
        from chatterbox.tts_turbo import ChatterboxTurboTTS
        model = ChatterboxTurboTTS.from_pretrained(device=device)
        # Fix: Turbo's loudness step returns 64-bit audio with current numpy, which crashes
        # the reference processing ("expected scalar type Double but found Float").
        orig = model.norm_loudness
        model.norm_loudness = lambda wav, sr: np.asarray(orig(wav, sr), dtype=np.float32)
    else:
        from chatterbox.tts import ChatterboxTTS
        model = ChatterboxTTS.from_pretrained(device=device)
    print(f"Model ready in {time.time() - t:.0f}s.")
    return model


def check_ref(path, kind):
    if not path.exists():
        sys.exit(f"Reference clip not found: {path}")
    info = sf.info(str(path))
    if info.duration < 5.5:
        sys.exit(f"{path.name} is {info.duration:.1f}s. Use a reference clip of 7-15 seconds.")
    if info.duration > 20:
        print(f"Note: {path.name} is {info.duration:.0f}s; only about the first 10-15s is used.")


def load_voice_profile(model, kind, path, device):
    """Load a voice profile made by build_voice.py (identity averaged over many clips)."""
    if not path.exists():
        sys.exit(f"Voice profile not found: {path}")
    meta = path.with_suffix(".json")
    if meta.exists():
        made_for = json.loads(meta.read_text(encoding="utf-8")).get("model", kind)
        if made_for != kind:
            sys.exit(f"{path.name} was built for the {made_for} model. "
                     f"Add --model {made_for}, or rebuild it with --model {kind}.")
    if kind == "turbo":
        from chatterbox.tts_turbo import Conditionals
    else:
        from chatterbox.tts import Conditionals
    model.conds = Conditionals.load(str(path), map_location="cpu").to(device)


def prepare_voice(model, kind, ref, exaggeration):
    if kind == "turbo":
        model.prepare_conditionals(str(ref), exaggeration=0.0, norm_loudness=True)
    else:
        model.prepare_conditionals(str(ref), exaggeration=exaggeration)


def generate_piece(model, kind, text, args):
    if kind == "turbo":
        wav = model.generate(text, temperature=args.temperature)
    else:
        wav = model.generate(text, exaggeration=args.exaggeration,
                             cfg_weight=args.cfg, temperature=args.temperature)
    return wav.squeeze().detach().cpu().numpy().astype(np.float32)


def seed_everything(seed):
    import torch
    random.seed(seed)
    np.random.seed(seed % (2 ** 32))
    torch.manual_seed(seed)


# ---------------------------------------------------------------- audio assembly

def trim(wav, sr, top_db=45, keep=0.04):
    """Trim silence at the ends of a piece, leaving a few ms so words aren't clipped."""
    import librosa
    if len(wav) < sr // 10:
        return wav
    _, (start, end) = librosa.effects.trim(wav, top_db=top_db)
    pad = int(sr * keep)
    return wav[max(0, start - pad): min(len(wav), end + pad)]


def assemble(pieces, sr, pause, para_pause):
    parts = []
    for wav, kind in pieces:
        parts.append(wav)
        gap = para_pause if kind == "paragraph" else pause
        parts.append(np.zeros(int(sr * gap), dtype=np.float32))
    audio = np.concatenate(parts[:-1]) if len(parts) > 1 else parts[0]
    peak = float(np.max(np.abs(audio))) if audio.size else 0.0
    if peak > 0:
        audio = audio * (0.89 / peak)  # about -1 dBFS, same level as the dataset clips
    return audio


# ---------------------------------------------------------------- main

def render(model, kind, text, out_path, args, seed):
    chunks = split_text(text, args.max_chars)
    if not chunks:
        print("Nothing to say (empty text).")
        return None
    seed_everything(seed)
    pieces, t0 = [], time.time()
    for i, (chunk, pause_kind) in enumerate(chunks, 1):
        t = time.time()
        print(f"  [{i}/{len(chunks)}] {chunk[:80]}{'...' if len(chunk) > 80 else ''}")
        wav = generate_piece(model, kind, chunk, args)
        pieces.append((trim(wav, model.sr), pause_kind))
        took = time.time() - t
        print(f"      {len(wav) / model.sr:.1f}s of audio in {took:.0f}s")
    audio = assemble(pieces, model.sr, args.pause, args.para_pause)

    out_path.parent.mkdir(parents=True, exist_ok=True)
    sf.write(str(out_path), audio, model.sr, subtype="PCM_16")
    settings = {
        "text": text, "reference": str(args.ref) if args.ref else None,
        "voice_profile": str(args.voice) if args.voice else None, "model": kind, "seed": seed,
        "exaggeration": args.exaggeration, "cfg": args.cfg, "temperature": args.temperature,
        "pause": args.pause, "para_pause": args.para_pause, "max_chars": args.max_chars,
        "created": datetime.now().isoformat(timespec="seconds"),
        "note": "AI-generated voice (Chatterbox, Perth watermark).",
    }
    out_path.with_suffix(".json").write_text(json.dumps(settings, indent=2), encoding="utf-8")
    print(f"Saved {out_path}  ({len(audio) / model.sr:.1f}s, made in {time.time() - t0:.0f}s, "
          f"seed {seed})")
    return out_path


def main():
    p = argparse.ArgumentParser(description="Generate speech in a cloned voice (Chatterbox).")
    p.add_argument("text", nargs="?", help="Text to speak (or use --file / --interactive)")
    p.add_argument("--file", type=Path, help="Text file with the script (blank line = paragraph)")
    p.add_argument("--interactive", action="store_true",
                   help="Keep the model loaded and type lines one after another")
    p.add_argument("--ref", type=Path, help="Reference clip of Brian (7-15 s WAV)")
    p.add_argument("--voice", type=Path,
                   help="Voice profile from build_voice.py (built from all clips), instead of --ref")
    p.add_argument("--out", type=Path, help="Output WAV (default: output\\brian_<time>.wav)")
    p.add_argument("--model", choices=["standard", "turbo"], default="standard",
                   help="standard: best control and quality. turbo: faster, fewer controls")
    p.add_argument("--exaggeration", type=float, default=0.5,
                   help="Emotion/energy, 0.25 calm to 1.0 very animated (standard model only)")
    p.add_argument("--cfg", type=float, default=0.5,
                   help="Pacing: lower (0.3) = slower, more deliberate; higher = faster (standard only)")
    p.add_argument("--temperature", type=float, default=0.8,
                   help="Variation between takes, 0.5 steady to 1.0 lively")
    p.add_argument("--seed", type=int, help="Repeat an exact take (seed is in each .json file)")
    p.add_argument("--takes", type=int, default=1, help="Make several versions to choose from")
    p.add_argument("--pause", type=float, default=0.25, help="Seconds between sentences")
    p.add_argument("--para-pause", type=float, default=0.7, help="Seconds between paragraphs")
    p.add_argument("--max-chars", type=int, default=280, help="Longest piece sent to the model")
    p.add_argument("--device", default="auto", help="auto, cpu, cuda")
    args = p.parse_args()

    if not (args.text or args.file or args.interactive):
        p.error("give some text, --file script.txt, or --interactive")
    if bool(args.ref) == bool(args.voice):
        p.error("use either --ref clip.wav or --voice profile.pt (one of them)")
    text = args.file.read_text(encoding="utf-8") if args.file else args.text
    if args.ref:
        check_ref(args.ref, args.model)
    else:
        if not args.voice.exists():
            sys.exit(f"Voice profile not found: {args.voice}")
        meta = args.voice.with_suffix(".json")
        if meta.exists():
            made_for = json.loads(meta.read_text(encoding="utf-8")).get("model", args.model)
            if made_for != args.model:
                sys.exit(f"{args.voice.name} was built for the {made_for} model. "
                         f"Add --model {made_for}, or rebuild it with --model {args.model}.")

    device = pick_device(args.device)
    if device == "cpu":
        print("Running on CPU: expect roughly 3-8 seconds of processing per second of speech.")
    model = load_model(args.model, device)
    if args.voice:
        print(f"Loading voice profile {args.voice.name}...")
        load_voice_profile(model, args.model, args.voice, device)
    else:
        print(f"Learning the voice from {args.ref.name}...")
        prepare_voice(model, args.model, args.ref, args.exaggeration)

    stamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_seed = args.seed if args.seed is not None else random.randint(0, 2 ** 31 - 1)

    if args.interactive:
        print("\nType a line and press Enter to generate it. Empty line to quit.")
        n = 0
        while True:
            try:
                line = input("\n> ").strip()
            except (EOFError, KeyboardInterrupt):
                break
            if not line:
                break
            n += 1
            render(model, args.model, line, Path("output") / f"line_{stamp}_{n:03d}.wav",
                   args, base_seed + n)
        return

    out = args.out or Path("output") / f"brian_{stamp}.wav"
    for k in range(args.takes):
        path = out if args.takes == 1 else out.with_name(f"{out.stem}_take{k + 1}{out.suffix}")
        if args.takes > 1:
            print(f"\nTake {k + 1} of {args.takes}")
        render(model, args.model, text, path, args, base_seed + k)


if __name__ == "__main__":
    main()