import torch
import subprocess
import os
import gradio as gr
import json
from datetime import datetime
import tempfile
from PIL import Image
import cv2
import numpy as np
from huggingface_hub import snapshot_download
import gc
import psutil

# =============================================================================
# قسم الإعدادات الأولية
# =============================================================================

print("\n--- بدء الإعدادات الأولية ---")

# التحقق من GPU المتاح ومساحة القرص
print("\n[1/4] التحقق من موارد النظام...")
if torch.cuda.is_available():
    gpu_name = torch.cuda.get_device_name(0)
    gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
    print(f"✅ GPU متاح: {gpu_name}")
    print(f"📊 ذاكرة GPU: {gpu_memory:.1f} GB")
    if gpu_memory < 15:
        print("⚠️ تحذير: قد تحتاج إلى GPU بذاكرة أكبر للحصول على أفضل أداء.")
else:
    print("❌ لا يوجد GPU متاح. يُنصح بتفعيل GPU في إعدادات Colab أو استخدام بيئة تدعم GPU.")

try:
    disk_usage = subprocess.check_output(["df", "-h", "/"]).decode("utf-8")
    print("\n💾 مساحة القرص:")
    print(disk_usage.split("\n")[1])
except Exception as e:
    print(f"❌ خطأ في التحقق من مساحة القرص: {e}")

# تحميل مستودع AccVideo AI
print("\n[3/4] تحميل مستودع AccVideo AI...")
accvideo_path = "/tmp/AccVideo"
if not os.path.exists(accvideo_path):
    try:
        subprocess.run(["git", "clone", "https://github.com/aejion/AccVideo.git", "/tmp/AccVideo"], check=True)
        print("✅ تم تحميل AccVideo بنجاح!")
    except Exception as e:
        print(f"❌ خطأ في تحميل AccVideo: {e}")
else:
    print("✅ AccVideo موجود بالفعل!")

# الانتقال إلى مجلد AccVideo وتثبيت متطلباته
print("\n[4/4] الانتقال إلى مجلد AccVideo وتثبيت متطلباته...")
try:
    os.chdir(accvideo_path)
    print(f"📁 المجلد الحالي: {os.getcwd()}")
    if os.path.exists("requirements.txt"):
        subprocess.run(["pip", "install", "-r", "requirements.txt"], check=True)
        print("✅ تم تثبيت متطلبات AccVideo!")
    else:
        print("⚠️ ملف requirements.txt غير موجود في مجلد AccVideo.")
except Exception as e:
    print(f"❌ خطأ في الانتقال أو تثبيت المتطلبات: {e}")

print("\n--- انتهت الإعدادات الأولية ---")

# =============================================================================
# وظائف تحميل النموذج وإنشاء الفيديو
# =============================================================================

# خيارات النماذج المتاحة
models_info = {
    "HunyuanT2V (الأساسي)": {
        "repo_id": "aejion/AccVideo",
        "local_dir": "./ckpts",
        "description": "نموذج أساسي لتحويل النص إلى فيديو"
    },
    "WanX-T2V-14B (متقدم)": {
        "repo_id": "aejion/AccVideo-WanX-T2V-14B",
        "local_dir": "./wanx_t2v_ckpts",
        "description": "نموذج متقدم عالي الجودة للنص إلى فيديو"
    },
    "WanX-I2V-480P-14B (صورة إلى فيديو)": {
        "repo_id": "aejion/AccVideo-WanX-I2V-480P-14B",
        "local_dir": "./wanx_i2v_ckpts",
        "description": "نموذج لتحويل الصورة إلى فيديو"
    }
}

# متغيرات عامة لحالة النموذج
current_model = None
model_type = None

def download_model_from_hf(selected_model_name):
    """تحميل النموذج المحدد من Hugging Face"""
    model_info = models_info[selected_model_name]
    print(f"🔄 جاري تحميل {selected_model_name}...")
    print(f"📝 الوصف: {model_info['description']}")
    
    try:
        if not os.path.exists(model_info["local_dir"]):
            snapshot_download(
                repo_id=model_info["repo_id"],
                local_dir=model_info["local_dir"],
                local_dir_use_symlinks=False
            )
            print(f"✅ تم تحميل {selected_model_name} بنجاح!")
        else:
            print(f"✅ {selected_model_name} موجود بالفعل!")
        return True
    except Exception as e:
        print(f"❌ خطأ في تحميل النموذج: {str(e)}")
        return False

def load_model(model_choice):
    """تحميل النموذج المحدد وتحديث الحالة العالمية"""
    global current_model, model_type
    
    if not download_model_from_hf(model_choice):
        return "❌ فشل تحميل النموذج. يرجى التحقق من الاتصال بالإنترنت أو مساحة القرص."

    try:
        if model_choice == "HunyuanT2V (الأساسي)":
            model_type = "hunyuan"
            return "✅ تم تحميل نموذج HunyuanT2V بنجاح!"
                
        elif model_choice == "WanX-T2V-14B (متقدم)":
            model_type = "wanx_t2v"
            return "✅ تم تحميل نموذج WanX-T2V-14B بنجاح!"
                
        elif model_choice == "WanX-I2V-480P-14B (صورة إلى فيديو)":
            model_type = "wanx_i2v"
            return "✅ تم تحميل نموذج WanX-I2V-480P-14B بنجاح!"
                
    except Exception as e:
        return f"❌ خطأ في إعداد النموذج: {str(e)}"

def generate_video(prompt, image_input, width, height, num_frames, steps, guidance_scale, seed):
    """إنشاء الفيديو بناءً على المدخلات ونوع النموذج المحدد"""
    global model_type
    
    if model_type is None:
        return None, "❌ يرجى تحميل النموذج أولاً من علامة تبويب 'إعداد النموذج'."
    
    try:
        # إنشاء مجلد النتائج
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        output_dir = os.path.join(os.getcwd(), "results", f"video_{timestamp}")
        os.makedirs(output_dir, exist_ok=True)
        
        # حفظ النص في ملف
        prompt_file = os.path.join(output_dir, "prompt.txt")
        with open(prompt_file, "w", encoding="utf-8") as f:
            f.write(prompt)
        
        cmd = []
        # تحديد الأمر حسب نوع النموذج
        if model_type == "hunyuan":
            cmd = [
                "python", "sample_t2v.py",
                "--height", str(height),
                "--width", str(width),
                "--num_frames", str(num_frames),
                "--num_inference_steps", str(steps),
                "--guidance_scale", "1",
                "--embedded_cfg_scale", str(guidance_scale),
                "--flow_shift", "7",
                "--flow-reverse",
                "--prompt_file", prompt_file,
                "--seed", str(seed),
                "--output_path", output_dir,
                "--model_path", "./ckpts",
                "--dit-weight", "./ckpts/accvideo-t2v-5-steps/diffusion_pytorch_model.pt"
            ]
            
        elif model_type == "wanx_t2v":
            cmd = [
                "python", "sample_wanx_t2v.py",
                "--task", "t2v-14B",
                "--size", f"{width}*{height}",
                "--ckpt_dir", "./wanx_t2v_ckpts",
                "--sample_solver", "unipc",
                "--save_dir", output_dir,
                "--sample_steps", str(steps),
                "--prompt", prompt
            ]
            
        elif model_type == "wanx_i2v":
            if image_input is None:
                return None, "❌ يرجى رفع صورة لاستخدام نموذج الصورة إلى فيديو."
                
            # حفظ الصورة المدخلة
            image_path = os.path.join(output_dir, "input_image.jpg")
            image_input.save(image_path)
            
            cmd = [
                "python", "sample_wanx_i2v.py",
                "--task", "i2v-14B",
                "--size", f"{width}*{height}",
                "--ckpt_dir", "./wanx_i2v_ckpts",
                "--sample_solver", "unipc",
                "--save_dir", output_dir,
                "--sample_steps", str(steps),
                "--image_path", image_path,
                "--prompt", prompt
            ]
        
        print(f'🔄 جاري إنشاء الفيديو... الأمر: {" ".join(cmd)}')
        
        # تشغيل الأمر في مجلد AccVideo
        result = subprocess.run(cmd, capture_output=True, text=True, cwd=os.getcwd())
        
        if result.returncode == 0:
            # البحث عن ملف الفيديو المُنتج
            video_files = []
            for root, dirs, files in os.walk(output_dir):
                for file in files:
                    if file.endswith((".mp4", ".avi", ".mov")):
                        video_files.append(os.path.join(root, file))
            
            if video_files:
                video_path = video_files[0]
                return video_path, f"✅ تم إنشاء الفيديو بنجاح!\n📁 المسار: {video_path}"
            else:
                return None, f"❌ لم يتم العثور على ملف فيديو في {output_dir}.\nالخرج: {result.stdout}\nالخطأ: {result.stderr}"
        else:
            error_msg = result.stderr if result.stderr else result.stdout
            return None, f"❌ خطأ في إنشاء الفيديو:\n{error_msg}"
            
    except Exception as e:
        return None, f"❌ خطأ غير متوقع أثناء إنشاء الفيديو: {str(e)}"

print("✅ تم إعداد وظائف الواجهة بنجاح!")

# =============================================================================
# واجهة المستخدم الرسومية (Gradio)
# =============================================================================

with gr.Blocks(title="AccVideo AI - مولد الفيديو بالذكاء الاصطناعي", theme=gr.themes.Soft()) as interface:
    
    gr.HTML("""
    <div style="text-align: center; padding: 20px;">
        <h1>🎬 AccVideo AI</h1>
        <h2>مولد الفيديو بالذكاء الاصطناعي</h2>
        <p>إنشاء مقاطع فيديو عالية الجودة من النصوص والصور</p>
    </div>
    """)
    
    with gr.Tab("🎯 إعداد النموذج"):
        with gr.Row():
            model_choice = gr.Dropdown(
                choices=list(models_info.keys()),
                value="HunyuanT2V (الأساسي)",
                label="اختر النموذج",
            )
            load_btn = gr.Button("تحميل النموذج", variant="primary")
        
        model_status = gr.Textbox(label="حالة النموذج", interactive=False)
        
        load_btn.click(
            fn=load_model,
            inputs=[model_choice],
            outputs=[model_status]
        )
    
    with gr.Tab("🎨 إنشاء الفيديو"):
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML("<h3>📝 الإدخال</h3>")
                
                prompt_input = gr.Textbox(
                    label="وصف الفيديو",
                    placeholder="اكتب وصفاً مفصلاً للفيديو الذي تريد إنشاءه...",
                    lines=3,
                )
                
                image_input = gr.Image(
                    label="الصورة المرجعية (اختياري)",
                    type="pil",
                )
                
                gr.HTML("<h3>⚙️ الإعدادات</h3>")
                
                with gr.Row():
                    width_input = gr.Slider(
                        minimum=256, maximum=1280, value=832, step=64,
                        label="العرض",
                    )
                    height_input = gr.Slider(
                        minimum=256, maximum=720, value=480, step=64,
                        label="الارتفاع",
                    )
                
                with gr.Row():
                    frames_input = gr.Slider(
                        minimum=16, maximum=120, value=81, step=1,
                        label="عدد الإطارات",
                    )
                    steps_input = gr.Slider(
                        minimum=5, maximum=20, value=10, step=1,
                        label="خطوات المعالجة",
                    )
                
                with gr.Row():
                    guidance_input = gr.Slider(
                        minimum=1.0, maximum=10.0, value=6.0, step=0.5,
                        label="قوة التوجيه",
                    )
                    seed_input = gr.Number(
                        value=1024, label="البذرة",
                    )
                
                generate_btn = gr.Button("🎬 إنشاء الفيديو", variant="primary", size="lg")
            
            with gr.Column(scale=1):
                gr.HTML("<h3>🎥 النتيجة</h3>")
                
                video_output = gr.Video(
                    label="الفيديو المُنتج",
                    height=400
                )
                
                status_output = gr.Textbox(
                    label="حالة المعالجة",
                    lines=5,
                    interactive=False
                )
        
        # ربط الأحداث
        generate_btn.click(
            fn=generate_video,
            inputs=[
                prompt_input, image_input, width_input, height_input,
                frames_input, steps_input, guidance_input, seed_input
            ],
            outputs=[video_output, status_output]
        )
    
    with gr.Tab("📚 المساعدة والنصائح"):
        gr.HTML("""
        <div style="text-align: center; padding: 20px;">
            <h2>💡 نصائح للحصول على أفضل النتائج</h2>
            
            <h3>📝 كتابة الوصف:</h3>
            <ul>
                <li>استخدم وصفاً واضحاً ومفصلاً</li>
                <li>اذكر التفاصيل المهمة مثل الألوان والحركة والإضاءة</li>
                <li>تجنب الوصف المعقد أو المتناقض</li>
                <li>مثال جيد: "قطة برتقالية تلعب بكرة صوف حمراء في حديقة خضراء مشمسة"</li>
            </ul>
            
            <h3>⚙️ الإعدادات:</h3>
            <ul>
                <li><strong>الأبعاد:</strong> ابدأ بأبعاد صغيرة للاختبار (832x480)</li>
                <li><strong>عدد الإطارات:</strong> 81 إطار ≈ 3 ثوانٍ، 120 إطار ≈ 4.5 ثانية</li>
                <li><strong>خطوات المعالجة:</strong> 10 خطوات توازن جيد بين الجودة والسرعة</li>
                <li><strong>قوة التوجيه:</strong> 6.0 قيمة جيدة للبداية</li>
            </ul>
            
            <h3>🚀 تحسين الأداء:</h3>
            <ul>
                <li>استخدم GPU عالي الأداء (T4 أو أفضل)</li>
                <li>تأكد من وجود مساحة كافية على القرص</li>
                <li>أغلق التطبيقات الأخرى لتوفير الذاكرة</li>
            </ul>
            
            <h3>🎯 النماذج المختلفة:</h3>
            <ul>
                <li><strong>HunyuanT2V:</strong> نموذج أساسي سريع</li>
                <li><strong>WanX-T2V-14B:</strong> جودة عالية، يحتاج ذاكرة أكبر</li>
                <li><strong>WanX-I2V-480P-14B:</strong> لتحويل الصور إلى فيديو</li>
            </ul>
            
            <h3>⏱️ أوقات المعالجة المتوقعة:</h3>
            <ul>
                <li>فيديو 832x480، 81 إطار: 2-5 دقائق</li>
                <li>فيديو 1280x720، 120 إطار: 10-15 دقيقة</li>
            </ul>
        </div>
        """)

# تشغيل الواجهة
if __name__ == "__main__":
    interface.launch(
        share=True,
        debug=True,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True
    )