from TTS.api import TTS
import torch
from pydub import AudioSegment
import random
import string
import os

# ✅ Thêm allowlist các class bị giới hạn nếu dùng PyTorch >2.6
torch.serialization.add_safe_globals([
    # Đảm bảo các class cần cho XTTS được cho phép
    __import__("TTS.tts.configs.xtts_config", fromlist=["XttsConfig"]).XttsConfig,
    __import__("TTS.tts.models.xtts", fromlist=["Xtts"]).Xtts,
    __import__("TTS.tts.models.xtts", fromlist=["XttsAudioConfig"]).XttsAudioConfig,
    __import__("TTS.config.shared_configs", fromlist=["BaseDatasetConfig"]).BaseDatasetConfig
])

# ✅ Load XTTSv2 model
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2", gpu=False)

# ✅ Hàm chính để sinh giọng nói từ văn bản + mẫu giọng
def generate_tts_xtts(text, speaker_wav, speed=1.0):
    output_path = "xtts_output_" + ''.join(random.choices(string.ascii_lowercase, k=6)) + ".wav"

    wav = tts.tts(
        text=text,
        speaker_wav=speaker_wav,
        language="en",  # 🔁 Chuyển "vi" nếu muốn clone giọng Việt
        speed=speed
    )

    tts.save_wav(wav, output_path)

    # ✅ Convert WAV to MP3
    mp3_path = output_path.replace(".wav", ".mp3")
    sound = AudioSegment.from_wav(output_path)
    sound.export(mp3_path, format="mp3")
    os.remove(output_path)

    return mp3_path