import gradio as gr
import cv2
import numpy as np
from PIL import Image
import tempfile
import os
import zipfile
import trimesh
from pathlib import Path

def extract_canny_edges(video_path, low_threshold=50, high_threshold=150):
    """
    استخراج Canny edges از ویدیو
    """
    cap = cv2.VideoCapture(video_path)
    
    # دریافت اطلاعات ویدیو
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    
    # ساخت فایل خروجی موقت برای ویدیو
    output_video_path = tempfile.mktemp(suffix='.mp4')
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
    
    frame_count = 0
    all_frames = []
    canny_frames_preview = []
    
    # ساخت پوشه موقت برای فریم‌ها
    frames_dir = tempfile.mkdtemp()
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        # تبدیل به Grayscale
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        
        # اعمال Gaussian Blur برای کاهش نویز
        blurred = cv2.GaussianBlur(gray, (5, 5), 0)
        
        # استخراج Canny edges
        edges = cv2.Canny(blurred, low_threshold, high_threshold)
        
        # تبدیل به BGR برای ذخیره در ویدیو
        edges_bgr = cv2.cvtColor(edges, cv2.COLOR_GRAY2BGR)
        
        # نوشتن فریم در ویدیوی خروجی
        out.write(edges_bgr)
        
        # ذخیره فریم به صورت تصویر جداگانه
        frame_filename = os.path.join(frames_dir, f'frame_{frame_count:05d}.png')
        cv2.imwrite(frame_filename, edges_bgr)
        
        # ذخیره برای پیش‌نمایش (هر 5 فریم)
        if frame_count % 5 == 0:
            edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
            canny_frames_preview.append(Image.fromarray(edges_rgb))
        
        # ذخیره برای تبدیل به 3D
        all_frames.append(edges)
        
        frame_count += 1
    
    cap.release()
    out.release()
    
    return output_video_path, frames_dir, all_frames, canny_frames_preview, frame_count, fps, width, height

def create_frames_zip(frames_dir):
    """
    ساخت فایل ZIP از تمام فریم‌ها
    """
    zip_path = tempfile.mktemp(suffix='.zip')
    
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(frames_dir):
            for file in sorted(files):
                if file.endswith('.png'):
                    file_path = os.path.join(root, file)
                    arcname = file
                    zipf.write(file_path, arcname)
    
    return zip_path

def create_3d_model(frames, width, height, depth_scale=0.05):
    """
    تبدیل فریم‌های Canny به مدل 3D (GLB) با Mesh واقعی
    """
    try:
        # انتخاب تعدادی از فریم‌ها برای جلوگیری از سنگین شدن مدل
        step = max(1, len(frames) // 20)  # حداکثر 20 لایه
        selected_frames = frames[::step]
        
        all_vertices = []
        all_faces = []
        vertex_count = 0
        
        # تبدیل هر فریم به یک لایه 3D
        for idx, frame in enumerate(selected_frames):
            # کاهش وضوح برای سبک‌تر کردن مدل
            scale = 8
            small_frame = cv2.resize(frame, (width // scale, height // scale))
            
            h, w = small_frame.shape
            z = idx * depth_scale  # فاصله بین لایه‌ها
            
            # ایجاد یک grid از vertices برای این فریم
            layer_vertices = []
            layer_indices = {}
            
            for y in range(h):
                for x in range(w):
                    if small_frame[y, x] > 128:  # پیکسل‌های روشن
                        # نرمال‌سازی مختصات
                        nx = (x / w - 0.5) * 2
                        ny = (0.5 - y / h) * 2
                        
                        vertex_idx = len(all_vertices)
                        all_vertices.append([nx, ny, z])
                        layer_vertices.append(vertex_idx)
                        layer_indices[(x, y)] = vertex_idx
            
            # ایجاد faces درون هر لایه (اتصال پیکسل‌های مجاور)
            for y in range(h - 1):
                for x in range(w - 1):
                    # بررسی اگر این پیکسل و همسایگانش سفید باشند
                    if ((x, y) in layer_indices and 
                        (x + 1, y) in layer_indices and 
                        (x, y + 1) in layer_indices):
                        
                        v1 = layer_indices[(x, y)]
                        v2 = layer_indices[(x + 1, y)]
                        v3 = layer_indices[(x, y + 1)]
                        
                        all_faces.append([v1, v2, v3])
                    
                    # مثلث دوم
                    if ((x + 1, y) in layer_indices and 
                        (x + 1, y + 1) in layer_indices and 
                        (x, y + 1) in layer_indices):
                        
                        v1 = layer_indices[(x + 1, y)]
                        v2 = layer_indices[(x + 1, y + 1)]
                        v3 = layer_indices[(x, y + 1)]
                        
                        all_faces.append([v1, v2, v3])
        
        if len(all_vertices) < 3 or len(all_faces) < 1:
            return None
        
        # ساخت Mesh واقعی
        vertices = np.array(all_vertices)
        faces = np.array(all_faces)
        
        # ایجاد mesh با trimesh
        mesh = trimesh.Trimesh(vertices=vertices, faces=faces)
        
        # اصلاح نرمال‌ها
        mesh.fix_normals()
        
        # ذخیره به فرمت GLB
        glb_path = tempfile.mktemp(suffix='.glb')
        mesh.export(glb_path, file_type='glb')
        
        return glb_path
        
    except Exception as e:
        print(f"خطا در ساخت مدل 3D: {e}")
        return None

def process_video(video_path, low_threshold, high_threshold, create_3d):
    """
    پردازش ویدیو و استخراج حرکات
    """
    if video_path is None:
        return None, None, None, None, "❌ لطفاً یک ویدیو آپلود کنید"
    
    try:
        # استخراج Canny edges
        output_video, frames_dir, all_frames, preview_frames, total_frames, fps, width, height = extract_canny_edges(
            video_path, 
            int(low_threshold), 
            int(high_threshold)
        )
        
        # ساخت ZIP از فریم‌ها
        frames_zip = create_frames_zip(frames_dir)
        
        # ساخت مدل 3D (اختیاری)
        glb_file = None
        if create_3d:
            glb_file = create_3d_model(all_frames, width, height)
            if glb_file is None:
                glb_status = "\n⚠️ ساخت مدل 3D با مشکل مواجه شد"
            else:
                glb_status = "\n✅ مدل 3D (GLB) آماده است"
        else:
            glb_status = ""
        
        status = f"""
        ✅ استخراج حرکات با موفقیت انجام شد!
        
        📊 اطلاعات:
        • تعداد کل فریم‌ها: {total_frames}
        • FPS: {fps}
        • مدت زمان: {total_frames/fps:.2f} ثانیه
        • رزولوشن: {width}x{height}
        {glb_status}
        
        📦 خروجی‌ها:
        • ویدیوی Canny: آماده برای دانلود
        • فریم‌های جداگانه: فایل ZIP شامل {total_frames} فریم
        {"• مدل 3D: فایل GLB برای نرم‌افزارهای سه‌بعدی" if glb_file else ""}
        """
        
        return output_video, frames_zip, glb_file, preview_frames, status
        
    except Exception as e:
        return None, None, None, None, f"❌ خطا: {str(e)}"

# رابط Gradio
with gr.Blocks(title="Wan2.1 Canny Edge Extractor Pro", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎬 استخراج حرکات ویدیو (Canny Edge Detection) - نسخه پیشرفته
    
    این ابزار با استفاده از الگوریتم **Canny Edge Detection**، لبه‌ها و حرکات ویدیو شما را استخراج کرده و در فرمت‌های مختلف ارائه می‌دهد:
    
    ✅ **ویدیوی کامل** | ✅ **فریم‌های جداگانه (ZIP)** | ✅ **مدل 3D (GLB)**
    """)
    
    with gr.Row():
        with gr.Column():
            input_video = gr.Video(
                label="📹 ویدیوی ورودی",
                height=400
            )
            
            gr.Markdown("### ⚙️ تنظیمات")
            
            low_threshold = gr.Slider(
                minimum=0,
                maximum=255,
                value=50,
                step=1,
                label="آستانه پایین (Low Threshold)",
                info="مقدار کمتر = لبه‌های بیشتر"
            )
            
            high_threshold = gr.Slider(
                minimum=0,
                maximum=255,
                value=150,
                step=1,
                label="آستانه بالا (High Threshold)",
                info="مقدار بیشتر = فقط لبه‌های قوی"
            )
            
            create_3d_checkbox = gr.Checkbox(
                label="🎲 ساخت مدل 3D (GLB)",
                value=True,
                info="فایل GLB برای Blender, Maya, Unity و..."
            )
            
            process_btn = gr.Button(
                "🚀 پردازش و استخراج",
                variant="primary",
                size="lg"
            )
        
        with gr.Column():
            status_text = gr.Textbox(
                label="وضعیت",
                lines=12,
                interactive=False
            )
            
            preview_gallery = gr.Gallery(
                label="🖼️ پیش‌نمایش فریم‌ها",
                columns=4,
                height=300
            )
    
    gr.Markdown("---")
    gr.Markdown("### 📥 دانلود خروجی‌ها")
    
    with gr.Row():
        with gr.Column():
            output_video = gr.Video(
                label="🎥 ویدیوی Canny Edges (کامل)",
                height=300
            )
        
        with gr.Column():
            frames_zip = gr.File(
                label="📦 فریم‌های جداگانه (ZIP)",
                file_types=[".zip"]
            )
        
        with gr.Column():
            glb_file = gr.File(
                label="🎲 مدل 3D (GLB)",
                file_types=[".glb"]
            )
    
    gr.Markdown("""
    ---
    ### 📖 راهنمای استفاده:
    
    **1. آپلود ویدیو**
    - ویدیوی خود را آپلود کنید (توصیه: 10-30 ثانیه برای نتیجه بهتر)
    
    **2. تنظیم پارامترها**
    - **آستانه پایین:** حساسیت تشخیص لبه (30-100 = جزئیات بیشتر)
    - **آستانه بالا:** فیلتر لبه‌های ضعیف (100-200 = فقط لبه‌های اصلی)
    - **مدل 3D:** برای استفاده در نرم‌افزارهای سه‌بعدی
    
    **3. دانلود خروجی‌ها**
    - 🎥 **ویدیو:** برای مشاهده یا استفاده در ControlNet
    - 📦 **ZIP:** تمام فریم‌ها به صورت PNG جداگانه
    - 🎲 **GLB:** برای Blender, Maya, Cinema 4D, Unity, Unreal Engine و...
    
    ---
    
    ### 🎯 کاربردها:
    
    **برای ControlNet:**
    - استفاده از ویدیو یا فریم‌ها به عنوان conditioning
    - کنترل دقیق حرکات در Text-to-Video
    
    **برای نرم‌افزارهای 3D:**
    - وارد کردن در Blender (File → Import → glTF)
    - استفاده در Maya, Cinema 4D, 3ds Max
    - پردازش در Unity یا Unreal Engine
    
    **برای Motion Graphics:**
    - استفاده از فریم‌های PNG در After Effects
    - ساخت انیمیشن فریم به فریم
    
    ---
    
    ### 💡 نکات مهم:
    
    ⚠️ **مدل 3D:**
    - فایل GLB یک Point Cloud از لبه‌هاست
    - هر فریم یک لایه در فضای Z می‌شود
    - برای مشاهده بهتر از نرم‌افزارهای 3D استفاده کنید
    
    ⚡ **بهینه‌سازی:**
    - ویدیوهای کوتاه‌تر = پردازش سریع‌تر
    - مدل 3D برای ویدیوهای طولانی ممکن است حجیم شود
    
    ---
    
    🔗 **مدل:** [TheDenk/wan2.1-t2v-1.3b-controlnet-canny-v1](https://huggingface.co/TheDenk/wan2.1-t2v-1.3b-controlnet-canny-v1)
    """)
    
    # اتصال دکمه به تابع
    process_btn.click(
        fn=process_video,
        inputs=[input_video, low_threshold, high_threshold, create_3d_checkbox],
        outputs=[output_video, frames_zip, glb_file, preview_gallery, status_text]
    )

# اجرای اپلیکیشن
if __name__ == "__main__":
    demo.launch(share=True)