import gc
import torch
import numpy as np
import gradio as gr
from transformers import VitsModel, AutoTokenizer

# Cache lưu trữ model gần nhất để tránh tải lại liên tục nếu dùng cùng 1 ngôn ngữ
current_model_id = None
loaded_model = None
loaded_tokenizer = None

def load_mms_model(lang_code):
    global current_model_id, loaded_model, loaded_tokenizer
    model_id = f"facebook/mms-tts-{lang_code.strip().lower()}"
    
    # Nếu chuyển sang ngôn ngữ mới, giải phóng model cũ khỏi bộ nhớ
    if current_model_id != model_id:
        if loaded_model is not None:
            del loaded_model
            del loaded_tokenizer
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
        
        print(f"Loading model: {model_id}...")
        loaded_tokenizer = AutoTokenizer.from_pretrained(model_id)
        loaded_model = VitsModel.from_pretrained(model_id)
        current_model_id = model_id
        
    return loaded_model, loaded_tokenizer

def generate_speech(text, selected_lang, custom_lang):
    if not text.strip():
        raise gr.Error("Vui lòng nhập văn bản!")
    
    # Ưu tiên lấy mã ngôn ngữ nhập tay nếu có, nếu không lấy từ danh sách chọn
    lang_code = custom_lang.strip() if custom_lang.strip() else selected_lang
    
    try:
        # Load model & tokenizer
        model, tokenizer = load_mms_model(lang_code)
        
        # Tiền xử lý văn bản
        inputs = tokenizer(text, return_tensors="pt")
        
        # Chạy dự đoán TTS
        with torch.no_grad():
            output = model(**inputs).waveform
            
        # Chuyển Tensor về Numpy Array chuẩn cho Gradio
        audio_data = output.squeeze().cpu().numpy()
        sampling_rate = model.config.sampling_rate
        
        return (sampling_rate, audio_data)
        
    except Exception as e:
        raise gr.Error(f"Lỗi khi tải ngôn ngữ '{lang_code}': {str(e)}. Hãy kiểm tra lại mã ngôn ngữ ISO-639-3.")

# Danh sách một số ngôn ngữ phổ biến làm mẫu
POPULAR_LANGUAGES = {
    "Tiếng Việt (vie)": "vie",
    "Tiếng Anh (eng)": "eng",
    "Tiếng Pháp (fra)": "fra",
    "Tiếng Tây Ban Nha (spa)": "spa",
    "Tiếng Trung (cmn)": "cmn",
    "Tiếng Nhật (jpn)": "jpn",
    "Tiếng Hàn (kor)": "kor",
    "Tiếng Thái (tha)": "tha",
    "Tiếng Lào (lao)": "lao",
    "Tiếng Khmer (khm)": "khm",
}

# Giao diện Gradio
with gr.Blocks(title="Facebook MMS TTS - Full Languages") as demo:
    gr.Markdown("# 🗣️ Facebook MMS TTS (Massively Multilingual Speech)")
    gr.Markdown(
        "Mô hình Chuyển đổi Văn bản thành Giọng nói hỗ trợ hơn 1.100+ ngôn ngữ từ Facebook (Meta)."
    )
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Nhập văn bản", 
                placeholder="Nhập đoạn văn bạn muốn chuyển thành giọng nói...",
                lines=4
            )
            
            lang_dropdown = gr.Dropdown(
                choices=list(POPULAR_LANGUAGES.values()),
                value="vie",
                label="Chọn ngôn ngữ phổ biến"
            )
            
            custom_lang_input = gr.Textbox(
                label="Hoặc nhập mã ngôn ngữ ISO-639-3 khác (Ví dụ: vie, eng, spa, ind...)",
                placeholder="Để trống nếu đã chọn ở danh sách trên..."
            )
            
            submit_btn = gr.Button("Tạo giọng nói 🔊", variant="primary")
            
        with gr.Column():
            audio_output = gr.Audio(label="Kết quả Âm thanh", type="numpy")

    submit_btn.click(
        fn=generate_speech,
        inputs=[input_text, lang_dropdown, custom_lang_input],
        outputs=[audio_output]
    )

if __name__ == "__main__":
    demo.launch()