import torchaudio as ta, tempfile, base64, os, torch
from huggingface_hub import snapshot_download
from safetensors.torch import load_file as load_safetensors
from chatterbox import mtl_tts

MODEL_NAME = "NAMAA-Space/NAMAA-Saudi-TTS"
device = "cuda" if torch.cuda.is_available() else "cpu"

ckpt_dir = snapshot_download(repo_id=MODEL_NAME)
model = mtl_tts.ChatterboxMultilingualTTS.from_pretrained(device=device)
t3_state = load_safetensors(f"{ckpt_dir}/t3_mtl23ls_v2.safetensors", device=device)
model.t3.load_state_dict(t3_state)
model.t3.to(device).eval()


def arTTS(text: str) -> str:
    with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tempFile:
        tempAudioPath = tempFile.name

    try:
        wav = model.generate(text, language_id="ar")
        ta.save(tempAudioPath, wav, model.sr)
        
        with open(tempAudioPath, "rb") as file:
            audioBase64 = base64.b64encode(file.read()).decode("utf-8")
    finally:
        os.remove(tempAudioPath) 

    return audioBase64