import spaces
import gradio as gr
import cv2
import numpy as np
from PIL import Image
import os
import shutil
import tempfile 
import datetime
import ffmpeg
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
# استيراد مكتبة السلايدر الجديدة
from gradio_imageslider import ImageSlider

# ==========================================
# 1. إعداد نموذج DDColor
# ==========================================
print("⏳ Loading DDColor Professional Model...")

try:
    ddcolor_pipeline = pipeline(
        Tasks.image_colorization, 
        model='damo/cv_ddcolor_image-colorization',
        device='gpu' 
    )
    print("✅ DDColor Model loaded successfully.")
except Exception as e:
    print(f"❌ Error loading DDColor model: {e}")
    ddcolor_pipeline = None

# ==========================================
# 2. دالة المعالجة
# ==========================================

@spaces.GPU(duration=180)
def colorize_video_professional(video_file):
    if not video_file:
        return None, None
    
    if ddcolor_pipeline is None:
        raise gr.Error("فشل تحميل النموذج.")

    print("🚀 Starting processing...")
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    temp_frames_dir = os.path.join(tempfile.gettempdir(), f"frames_{timestamp}")
    os.makedirs(temp_frames_dir, exist_ok=True)
    
    final_output_name = f"colored_ddcolor_{timestamp}.mp4"
    audio_path = os.path.join(tempfile.gettempdir(), f"audio_{timestamp}.aac")
    
    # مسارات صور المقارنة
    comp_original_path = os.path.join(tempfile.gettempdir(), f"comp_orig_{timestamp}.png")
    comp_colored_path = os.path.join(tempfile.gettempdir(), f"comp_color_{timestamp}.png")
    comparison_result = None

    # --- استخراج الصوت ---
    audio_exists = False
    try:
        ffmpeg.input(video_file).output(audio_path, acodec='copy').run(overwrite_output=True, quiet=True)
        audio_exists = True
    except ffmpeg.Error:
        pass

    # --- المعالجة ---
    cap = cv2.VideoCapture(video_file)
    fps = cap.get(cv2.CAP_PROP_FPS) or 25
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    
    # نختار إطاراً في المنتصف ليكون صورة المقارنة
    comparison_frame_index = max(0, total_frames // 2)
    
    frame_count = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
            
        # حفظ الإطار الأصلي للمقارنة إذا وصلنا للمنتصف
        if frame_count == comparison_frame_index:
            cv2.imwrite(comp_original_path, frame)
        
        # التلوين
        result = ddcolor_pipeline(frame)
        colorized_frame_bgr = result['output_img']
        
        # حفظ الإطار الملون للمقارنة
        if frame_count == comparison_frame_index:
            cv2.imwrite(comp_colored_path, colorized_frame_bgr)
            comparison_result = (comp_original_path, comp_colored_path)

        # حفظ الإطار للفيديو
        frame_filename = os.path.join(temp_frames_dir, f"frame_{frame_count:05d}.png")
        cv2.imwrite(frame_filename, colorized_frame_bgr)
        frame_count += 1
    
    cap.release()

    # --- تجميع الفيديو ---
    input_frames = ffmpeg.input(os.path.join(temp_frames_dir, 'frame_%05d.png'), framerate=fps)
    if audio_exists:
        stream = ffmpeg.output(input_frames, ffmpeg.input(audio_path), final_output_name, vcodec='libx264', pix_fmt='yuv420p', acodec='aac', shortest=None)
    else:
        stream = ffmpeg.output(input_frames, final_output_name, vcodec='libx264', pix_fmt='yuv420p')
    
    try:
        stream.run(overwrite_output=True, quiet=True)
    except ffmpeg.Error:
         # محاولة بديلة
        ffmpeg.input(os.path.join(temp_frames_dir, 'frame_%05d.png'), framerate=fps).output(final_output_name, vcodec='libx264', pix_fmt='yuv420p').run(overwrite_output=True)

    shutil.rmtree(temp_frames_dir, ignore_errors=True)
    
    # نرجع الفيديو + صور المقارنة (Tuple)
    return final_output_name, comparison_result

# ==========================================
# 3. الواجهة الجديدة
# ==========================================
custom_css = """
#col-container {max-width: 800px; margin-left: auto; margin-right: auto;}
"""

with gr.Blocks(css=custom_css, title="Pro Video Colorizer") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# 🎞️ Professional Video Colorizer (DDColor)")
        gr.Markdown("قم برفع فيديو أبيض وأسود وسيقوم الذكاء الاصطناعي بتلوينه.")
        
        with gr.Row():
            video_input = gr.Video(label="فيديو أبيض وأسود (Input)")
        
        submit_btn = gr.Button("✨ تلوين ومعاينة (Colorize)", variant="primary", size="lg")
        
        # عنصر المقارنة الجديد
        gr.Markdown("### 🔍 معاينة قبل وبعد (Before / After)")
        slider_output = ImageSlider(label="مقارنة النتيجة", type="filepath", position=0.5)
        
        gr.Markdown("### 🎥 الفيديو النهائي الملون")
        video_output = gr.Video(label="النتيجة النهائية")
        
        submit_btn.click(
            fn=colorize_video_professional,
            inputs=[video_input],
            outputs=[video_output, slider_output]
        )

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)