import gradio as gr
import cv2
import numpy as np
from PIL import Image
import tempfile
import os

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_path = tempfile.mktemp(suffix='.mp4')
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    out = cv2.VideoWriter(output_path, fourcc, fps, (width, height))
    
    frame_count = 0
    canny_frames = []
    
    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)
        
        # ذخیره برای پیش‌نمایش
        if frame_count % 5 == 0:  # هر 5 فریم یکی
            edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
            canny_frames.append(Image.fromarray(edges_rgb))
        
        frame_count += 1
    
    cap.release()
    out.release()
    
    return output_path, canny_frames, frame_count, fps

def process_video(video_path, low_threshold, high_threshold):
    """
    پردازش ویدیو و استخراج حرکات
    """
    if video_path is None:
        return None, None, "❌ لطفاً یک ویدیو آپلود کنید"
    
    try:
        output_video, preview_frames, total_frames, fps = extract_canny_edges(
            video_path, 
            int(low_threshold), 
            int(high_threshold)
        )
        
        status = f"""
        ✅ استخراج حرکات با موفقیت انجام شد!
        
        📊 اطلاعات:
        • تعداد کل فریم‌ها: {total_frames}
        • FPS: {fps}
        • مدت زمان: {total_frames/fps:.2f} ثانیه
        """
        
        return output_video, preview_frames, status
        
    except Exception as e:
        return None, None, f"❌ خطا: {str(e)}"

# رابط Gradio
with gr.Blocks(title="Wan2.1 Canny Edge Extractor", theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎬 استخراج حرکات ویدیو (Canny Edge Detection)
    
    این ابزار با استفاده از الگوریتم **Canny Edge Detection**، لبه‌ها و حرکات ویدیو شما را استخراج می‌کند.
    
    این خروجی می‌تواند به عنوان ورودی برای مدل **Wan2.1 ControlNet** استفاده شود.
    """)
    
    with gr.Row():
        with gr.Column():
            input_video = gr.Video(
                label="📹 ویدیوی ورودی",
                height=400
            )
            
            gr.Markdown("### ⚙️ تنظیمات Canny")
            
            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="مقدار بیشتر = فقط لبه‌های قوی"
            )
            
            process_btn = gr.Button(
                "🚀 استخراج حرکات",
                variant="primary",
                size="lg"
            )
        
        with gr.Column():
            status_text = gr.Textbox(
                label="وضعیت",
                lines=6,
                interactive=False
            )
            
            output_video = gr.Video(
                label="🎥 ویدیوی Canny Edges",
                height=400
            )
            
            preview_gallery = gr.Gallery(
                label="🖼️ پیش‌نمایش فریم‌ها",
                columns=4,
                height=300
            )
    
    gr.Markdown("""
    ---
    ### 📖 راهنمای استفاده:
    
    1. **آپلود ویدیو**: یک ویدیو آپلود کنید (توصیه: کمتر از 30 ثانیه)
    2. **تنظیم آستانه‌ها**: 
       - آستانه پایین: کنترل حساسیت تشخیص لبه
       - آستانه بالا: فیلتر کردن لبه‌های ضعیف
    3. **استخراج**: روی دکمه "استخراج حرکات" کلیک کنید
    4. **دانلود**: ویدیوی خروجی را دانلود کنید
    
    ### 💡 نکات:
    - **آستانه کم** (مثلاً 30-100): جزئیات بیشتر، نویز بیشتر
    - **آستانه متوسط** (مثلاً 50-150): پیشنهادی - تعادل خوب
    - **آستانه بالا** (مثلاً 100-200): فقط لبه‌های اصلی
    
    ### 🔗 استفاده در Wan2.1:
    ویدیوی خروجی می‌تواند به عنوان **ControlNet conditioning** برای مدل Wan2.1 استفاده شود.
    
    ---
    
    🔗 مدل: [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],
        outputs=[output_video, preview_gallery, status_text]
    )

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