import gradio as gr
from transformers import VitsModel, AutoTokenizer
import torch
import scipy.io.wavfile

# تحميل النموذج و "tokenizer"
model_name = "wasmdashai/vits-ar-sa-huba-v2"
try:
    model = VitsModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model_loaded_successfully = True
    print(f"تم تحميل النموذج '{model_name}' بنجاح.")
except Exception as e:
    model_loaded_successfully = False
    error_message = f"حدث خطأ أثناء تحميل النموذج '{model_name}': {e}"
    print(error_message)
    # يمكنك اختيار إيقاف التطبيق هنا أو عرض رسالة خطأ في الواجهة

# دالة التنبؤ (تحويل النص إلى كلام)
def text_to_speech(text):
    if not model_loaded_successfully:
        return None, f"خطأ: لم يتم تحميل النموذج. {error_message}"
    if not text:
        return None, "الرجاء إدخال نص."
    try:
        inputs = tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            output = model(**inputs).waveform

        # التأكد من أن المخرجات هي 1D tensor
        if output.ndim > 1:
            # إذا كانت متعددة الأبعاد، جرب أخذ القناة الأولى أو المتوسط
            # هنا نفترض أن القناة الأولى هي الصوت المطلوب
            waveform = output[0].cpu().numpy()
        else:
            waveform = output.cpu().numpy()

        # معدل أخذ العينات للنموذج (عادة ما يكون مذكورًا في بطاقة النموذج أو تكوينه)
        # إذا لم يكن متاحًا بشكل مباشر، قد تحتاج إلى التحقق من config النموذج
        # القيمة الشائعة لنماذج VITS هي 22050 أو 16000
        # بالنظر إلى بطاقة النموذج، يبدو أنها لا تحدد صراحةً، سنجرب 22050
        sampling_rate = model.config.sampling_rate if hasattr(model.config, 'sampling_rate') else 22050

        # حفظ الملف الصوتي مؤقتًا
        output_path = "output_audio.wav"
        scipy.io.wavfile.write(output_path, rate=sampling_rate, data=waveform)
        return output_path, None  # لا يوجد رسالة خطأ
    except Exception as e:
        print(f"حدث خطأ أثناء تحويل النص إلى كلام: {e}")
        return None, f"حدث خطأ أثناء المعالجة: {str(e)}"

# إنشاء واجهة Gradio
iface = gr.Interface(
    fn=text_to_speech,
    inputs=gr.Textbox(lines=3, placeholder="أدخل النص العربي هنا...", label="النص العربي"),
    outputs=[
        gr.Audio(label="الصوت الناتج", type="filepath"),
        gr.Textbox(label="رسالة الخطأ (إن وجدت)") # لإظهار رسائل الخطأ بشكل واضح
    ],
    title="تحويل النص إلى كلام باللغة العربية (نموذج vits-ar-sa-huda-v2)",
    description="أدخل نصًا باللغة العربية ليتم تحويله إلى كلام باستخدام نموذج VITS المدرب على اللهجة السعودية . النموذج من تطوير `wasmdashai`.",
    allow_flagging="never",
    examples=[
        ["السلام عليكم ورحمة الله وبركاته."],
        ["كيف حالك اليوم؟"],
        ["أهلاً وسهلاً بكم في هذه المساحة."],
    ]
)

# تشغيل الواجهة
if __name__ == "__main__":
    if model_loaded_successfully:
        iface.launch()
    else:
        # إذا لم يتم تحميل النموذج، قم بتشغيل واجهة بسيطة تعرض الخطأ
        def error_interface():
            with gr.Blocks() as demo:
                gr.Markdown(f"# خطأ في تحميل النموذج\n{error_message}")
            return demo
        error_interface().launch()