"""
Generic MOSS-TTS zero-shot cloning CLI: clone the voice in --reference and
say --text in it. Used directly, and as the final stage of webui/app.py.
"""

import argparse
import importlib.util

import soundfile as sf
import torch
from transformers import AutoModel, AutoProcessor

PRETRAINED = "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"


def resolve_attn_implementation(device: str, dtype: torch.dtype) -> str:
    if (
        device == "cuda"
        and importlib.util.find_spec("flash_attn") is not None
        and dtype in {torch.float16, torch.bfloat16}
    ):
        major, _ = torch.cuda.get_device_capability()
        if major >= 8:
            return "flash_attention_2"
    return "sdpa" if device == "cuda" else "eager"


def main() -> None:
    parser = argparse.ArgumentParser(description="MOSS-TTS zero-shot voice cloning")
    parser.add_argument("--text", required=True, help="Text to synthesize")
    parser.add_argument("--reference", required=True, help="Reference audio path")
    parser.add_argument("--output", required=True, help="Output wav path")
    parser.add_argument("--language", default="English", help="Language tag")
    parser.add_argument(
        "--instruction", default=None,
        help='Style/delivery instruction (e.g. "Change the emotion to excited."). '
        "Undocumented MOSS-TTS field for this checkpoint - confirmed to work for "
        "emotion in testing, passed through as-is.",
    )
    parser.add_argument(
        "--seed", type=int, default=None,
        help="Random seed for reproducibility. Unset = nondeterministic.",
    )
    parser.add_argument("--audio_temperature", type=float, default=1.7)
    parser.add_argument("--audio_top_p", type=float, default=0.8)
    parser.add_argument("--audio_top_k", type=int, default=25)
    parser.add_argument("--audio_repetition_penalty", type=float, default=1.0)
    parser.add_argument("--max_new_tokens", type=int, default=4096)
    args = parser.parse_args()

    device = "cuda" if torch.cuda.is_available() else "cpu"
    dtype = torch.bfloat16 if device == "cuda" else torch.float32

    if args.seed is not None:
        torch.manual_seed(args.seed)
        if device == "cuda":
            torch.cuda.manual_seed_all(args.seed)

    torch.backends.cuda.enable_cudnn_sdp(False)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    torch.backends.cuda.enable_math_sdp(True)

    attn_implementation = resolve_attn_implementation(device, dtype)
    print(f"[INFO] attn_implementation={attn_implementation}")

    processor = AutoProcessor.from_pretrained(PRETRAINED, trust_remote_code=True)
    processor.audio_tokenizer = processor.audio_tokenizer.to(device)

    model = AutoModel.from_pretrained(
        PRETRAINED,
        trust_remote_code=True,
        attn_implementation=attn_implementation,
        dtype=dtype,
    ).to(device)
    model.eval()

    conversation = [
        processor.build_user_message(
            text=args.text,
            language=args.language,
            reference=[args.reference],
            instruction=args.instruction,
        )
    ]

    with torch.no_grad():
        batch = processor([conversation], mode="generation")
        outputs = model.generate(
            input_ids=batch["input_ids"].to(device),
            attention_mask=batch["attention_mask"].to(device),
            max_new_tokens=args.max_new_tokens,
            do_sample=True,
            audio_temperature=args.audio_temperature,
            audio_top_p=args.audio_top_p,
            audio_top_k=args.audio_top_k,
            audio_repetition_penalty=args.audio_repetition_penalty,
        )
        message = processor.decode(outputs)[0]
        audio = message.audio_codes_list[0]
        data = audio.detach().cpu().to(torch.float32).numpy()
        if data.ndim == 2:
            data = data.T
        sf.write(args.output, data, processor.model_config.sampling_rate)
        print(f"[INFO] Saved {args.output}")


if __name__ == "__main__":
    main()