import os
import asyncio
import tempfile
import gradio as gr
import whisper
import torch
import edge_tts
from pathlib import Path
from moviepy.editor import VideoFileClip
from transformers import MarianMTModel, MarianTokenizer

# تهيئة النماذج
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
whisper_model = whisper.load_model("base")

# قاموس للغات المدعومة
SUPPORTED_LANGUAGES = {
    "ar": {"name": "العربية", "code": "ar-SA"},
    "en": {"name": "English", "code": "en-US"},
    "fr": {"name": "Français", "code": "fr-FR"},
    "es": {"name": "Español", "code": "es-ES"}
}

# قاموس لنماذج الترجمة
TRANSLATION_MODELS = {
    "ar-en": "Helsinki-NLP/opus-mt-ar-en",
    "en-ar": "Helsinki-NLP/opus-mt-en-ar",
    "fr-en": "Helsinki-NLP/opus-mt-fr-en",
    "en-fr": "Helsinki-NLP/opus-mt-en-fr",
    "es-en": "Helsinki-NLP/opus-mt-es-en",
    "en-es": "Helsinki-NLP/opus-mt-en-es"
}

# قاموس لأنواع الأصوات
VOICE_TYPES = {
    "رجل": {
        "ar": "ar-SA-HamedNeural",
        "en": "en-US-ChristopherNeural",
        "fr": "fr-FR-HenriNeural",
        "es": "es-ES-AlvaroNeural"
    },
    "امرأة": {
        "ar": "ar-SA-ZariyahNeural",
        "en": "en-US-JennyNeural",
        "fr": "fr-FR-DeniseNeural",
        "es": "es-ES-ElviraNeural"
    },
    "طفل": {
        "ar": "ar-SA-ZariyahNeural",
        "en": "en-US-JennyNeural",
        "fr": "fr-FR-DeniseNeural",
        "es": "es-ES-ElviraNeural"
    }
}

def extract_audio_from_video(video_path):
    """استخراج الصوت من الفيديو"""
    try:
        video = VideoFileClip(video_path)
        temp_audio_path = tempfile.mktemp(suffix=".mp3")
        video.audio.write_audiofile(temp_audio_path, codec='mp3')
        video.close()
        return temp_audio_path
    except Exception as e:
        raise Exception(f"خطأ في استخراج الصوت من الفيديو: {str(e)}")

def process_media_file(file_path, source_lang):
    """معالجة ملف الوسائط (صوت أو فيديو)"""
    try:
        if file_path is None:
            return "الرجاء تحميل ملف صوتي أو فيديو"
            
        # التحقق من نوع الملف
        if file_path.name.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            # إذا كان فيديو، استخرج الصوت منه
            audio_path = extract_audio_from_video(file_path.name)
        elif file_path.name.lower().endswith(('.mp3', '.wav', '.m4a')):
            # إذا كان ملف صوتي، استخدمه مباشرة
            audio_path = file_path.name
        else:
            return "نوع الملف غير مدعوم. الأنواع المدعومة هي: MP3, WAV, M4A, MP4, AVI, MOV, MKV"
        
        # تحويل الصوت إلى نص
        result = whisper_model.transcribe(audio_path, language=source_lang)
        
        # حذف الملف المؤقت إذا كان فيديو
        if file_path.name.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
            os.remove(audio_path)
        
        return result["text"]
    except Exception as e:
        return f"خطأ في معالجة الملف: {str(e)}"

def translate_text(text, source_lang, target_lang):
    """ترجمة النص باستخدام MarianMT"""
    if source_lang == target_lang:
        return text
    
    try:
        # إذا كانت اللغة المصدر ليست الإنجليزية، نترجم أولاً إلى الإنجليزية
        if source_lang != "en":
            model_name = TRANSLATION_MODELS[f"{source_lang}-en"]
            model = MarianMTModel.from_pretrained(model_name)
            tokenizer = MarianTokenizer.from_pretrained(model_name)
            
            inputs = tokenizer(text, return_tensors="pt", padding=True)
            translated = model.generate(**inputs)
            text = tokenizer.decode(translated[0], skip_special_tokens=True)
            
            # إذا كانت اللغة الهدف هي الإنجليزية، نتوقف هنا
            if target_lang == "en":
                return text
        
        # ترجمة من الإنجليزية إلى اللغة الهدف
        if target_lang != "en":
            model_name = TRANSLATION_MODELS[f"en-{target_lang}"]
            model = MarianMTModel.from_pretrained(model_name)
            tokenizer = MarianTokenizer.from_pretrained(model_name)
            
            inputs = tokenizer(text, return_tensors="pt", padding=True)
            translated = model.generate(**inputs)
            text = tokenizer.decode(translated[0], skip_special_tokens=True)
        
        return text
    except Exception as e:
        return f"خطأ في الترجمة: {str(e)}"

async def text_to_speech(text, language, voice_type):
    """تحويل النص إلى صوت باستخدام Edge TTS"""
    try:
        # إنشاء مجلد مؤقت للملفات الصوتية
        temp_dir = Path("temp_audio")
        temp_dir.mkdir(exist_ok=True)
        
        # اختيار الصوت المناسب
        voice = VOICE_TYPES[voice_type][language]
        
        # تعديل السرعة والنبرة حسب نوع الصوت
        rate = "+0%" if voice_type != "طفل" else "+15%"
        pitch = "+0Hz" if voice_type == "رجل" else "+10Hz" if voice_type == "امرأة" else "+60Hz"
        
        # إنشاء ملف صوتي مؤقت
        output_file = temp_dir / f"output_{voice_type}_{language}.mp3"
        
        # تكوين كائن communicate
        communicate = edge_tts.Communicate(text, voice, rate=rate, pitch=pitch)
        
        # حفظ الملف الصوتي
        await communicate.save(str(output_file))
        
        return str(output_file)
    
    except Exception as e:
        return f"خطأ في تحويل النص إلى صوت: {str(e)}"

def text_to_speech_wrapper(text, language, voice_type):
    """wrapper function لتشغيل الدالة غير المتزامنة"""
    return asyncio.run(text_to_speech(text, language, voice_type))

# إنشاء واجهة Gradio
with gr.Blocks(title="معالج الصوت والترجمة", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# معالج الصوت والترجمة متعدد اللغات")
    
    with gr.Tab("تحويل الوسائط إلى نص"):
        gr.Markdown("""
        ### الملفات المدعومة:
        - ملفات الصوت: MP3, WAV, M4A
        - ملفات الفيديو: MP4, AVI, MOV, MKV
        """)
        
        with gr.Row():
            media_input = gr.File(
                label="ملف صوتي أو فيديو"
            )
            source_lang = gr.Dropdown(
                choices=list(SUPPORTED_LANGUAGES.keys()),
                value="ar",
                label="لغة الملف"
            )
        
        transcribe_btn = gr.Button("تحويل إلى نص")
        transcribed_text = gr.Textbox(label="النص المستخرج", lines=5)
        
        transcribe_btn.click(
            fn=process_media_file,
            inputs=[media_input, source_lang],
            outputs=transcribed_text
        )
    
    with gr.Tab("ترجمة النص"):
        with gr.Row():
            input_text = gr.Textbox(label="النص المراد ترجمته", lines=5)
            translated_text = gr.Textbox(label="النص المترجم", lines=5)
        
        with gr.Row():
            trans_source_lang = gr.Dropdown(
                choices=list(SUPPORTED_LANGUAGES.keys()),
                value="ar",
                label="اللغة المصدر"
            )
            trans_target_lang = gr.Dropdown(
                choices=list(SUPPORTED_LANGUAGES.keys()),
                value="en",
                label="اللغة الهدف"
            )
        
        translate_btn = gr.Button("ترجمة")
        
        translate_btn.click(
            fn=translate_text,
            inputs=[input_text, trans_source_lang, trans_target_lang],
            outputs=translated_text
        )
    
    with gr.Tab("تحويل النص إلى صوت"):
        with gr.Row():
            tts_text = gr.Textbox(label="النص المراد تحويله إلى صوت", lines=5)
            tts_output = gr.Audio(label="الصوت الناتج")
        
        with gr.Row():
            tts_lang = gr.Dropdown(
                choices=list(SUPPORTED_LANGUAGES.keys()),
                value="ar",
                label="لغة النص"
            )
            voice_type = gr.Radio(
                choices=list(VOICE_TYPES.keys()),
                value="رجل",
                label="نوع الصوت"
            )
        
        tts_btn = gr.Button("تحويل إلى صوت")
        
        tts_btn.click(
            fn=text_to_speech_wrapper,
            inputs=[tts_text, tts_lang, voice_type],
            outputs=tts_output
        )

# تشغيل التطبيق
if __name__ == "__main__":
    demo.launch()