"""
Voice-clone + translate demo.

Takes a short English sentence spoken by the narrator in our sample clip
(the same sentence used as the reference audio, for an apples-to-apples
comparison), machine-translates it to Spanish, then uses MOSS-TTS to
synthesize the Spanish text in the narrator's cloned voice.

Reference audio (tts/reference/narrator_ref_en.wav) was trimmed from
demucs's isolated vocals stem for this same clip.
"""

import importlib.util
from pathlib import Path

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

HERE = Path(__file__).parent
REFERENCE_AUDIO = HERE / "reference" / "narrator_ref_en.wav"
OUTPUT_DIR = HERE / "output"

TEXT_EN = (
    "We shall show you in this film typical examples of children who have "
    "been well brought up, the kind who grow up to be strong, healthy, and "
    "happy like these young folks."
)

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


def translate_to_spanish(text: str) -> str:
    model_name = "Helsinki-NLP/opus-mt-en-es"
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name).to(device)
    batch = tokenizer([text], return_tensors="pt", padding=True).to(device)
    generated = model.generate(**batch, max_new_tokens=256)
    return tokenizer.decode(generated[0], skip_special_tokens=True)


def resolve_attn_implementation() -> 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:
    text_es = translate_to_spanish(TEXT_EN)
    print(f"[EN] {TEXT_EN}")
    print(f"[ES] {text_es}")

    # Disable the broken cuDNN SDPA backend (per MOSS-TTS README).
    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)

    # The flagship MOSS-TTS-v1.5 (8B) doesn't fit alongside its audio
    # tokenizer + generation overhead on a 24GB 3090 (OOMs at ~23.3GB used).
    # The 4B Local-Transformer variant is the recommended smaller-GPU option
    # and also natively outputs 48kHz stereo.
    pretrained = "OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5"
    attn_implementation = resolve_attn_implementation()
    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=text_es,
            language="Spanish",
            reference=[str(REFERENCE_AUDIO)],
        )
    ]

    OUTPUT_DIR.mkdir(exist_ok=True, parents=True)
    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=4096,
            do_sample=True,
            audio_temperature=1.7,
            audio_top_p=0.8,
            audio_top_k=25,
            audio_repetition_penalty=1.0,
        )
        message = processor.decode(outputs)[0]
        audio = message.audio_codes_list[0]
        # torchaudio.save() depends on torchcodec, whose bundled shared libs
        # failed to load in this environment; soundfile (already a MOSS-TTS
        # dependency) sidesteps that entirely.
        data = audio.detach().cpu().to(torch.float32).numpy()
        if data.ndim == 2:
            data = data.T  # soundfile wants (frames, channels), not (channels, frames)
        out_path = OUTPUT_DIR / "narrator_cloned_es.wav"
        sf.write(str(out_path), data, processor.model_config.sampling_rate)
        print(f"[INFO] Saved {out_path}")


if __name__ == "__main__":
    main()