#!/usr/bin/env python3
"""Generate one zero-shot voice-cloned utterance with the local NVFP4 bundle."""

from __future__ import annotations

import argparse
import re
from pathlib import Path

import soundfile as sf
import torch
from huggingface_hub import snapshot_download

from chatterbox_flash.nvfp4 import (
    BASE_ASSET_FILENAMES,
    BASE_CHECKPOINT_REVISION,
    BASE_REPO_ID,
    load_prepacked_tts_local,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--prompt-wav", type=Path, required=True)
    parser.add_argument("--text", required=True)
    parser.add_argument("--output", type=Path, default=Path("output.wav"))
    parser.add_argument(
        "--block-size",
        type=int,
        choices=(16, 24),
        default=16,
        help="16 is the quality profile; 24 is the optional faster profile",
    )
    parser.add_argument("--seed", type=int, default=4242)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if not args.prompt_wav.is_file():
        raise SystemExit(f"Missing reference audio: {args.prompt_wav}")

    bundle = Path(__file__).resolve().parent
    base_dir = Path(snapshot_download(
        repo_id=BASE_REPO_ID,
        revision=BASE_CHECKPOINT_REVISION,
        allow_patterns=list(BASE_ASSET_FILENAMES),
    ))
    capability = torch.cuda.get_device_capability()
    gpu_slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", torch.cuda.get_device_name())
    autotune_cache = (
        Path.home()
        / ".cache"
        / "chatterbox-flash-nvfp4"
        / f"flashinfer-{gpu_slug}-sm{capability[0]}{capability[1]}-d{args.block_size}.json"
    )

    torch.manual_seed(args.seed)
    torch.cuda.manual_seed_all(args.seed)
    tts = load_prepacked_tts_local(
        bundle,
        base_dir,
        device="cuda",
        drf_block_size=args.block_size,
        autotune_cache=autotune_cache,
    )
    waveform = tts.generate(
        args.text,
        audio_prompt_path=args.prompt_wav,
        exaggeration=0.5,
        num_steps=10,
        temperature=0.2,
        time_shift_tau=0.5,
        omnivoice_schedule_t_shift=0.5,
        cfg_scale=1.0,
        position_temperature=5.0,
        n_cfm_timesteps=2,
        backend="flashinfer",
        use_cuda_graph=True,
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    sf.write(args.output, waveform.float().numpy(), tts.sr)
    print(f"Wrote {args.output} at {tts.sr} Hz")


if __name__ == "__main__":
    main()