import gradio as gr
import numpy as np
from PIL import Image

def get_mask_bbox(mask: Image.Image):
    """احصل على إطار القناع (Bounding Box) كـ (xmin, ymin, xmax, ymax)"""
    mask_np = np.array(mask)
    if len(mask_np.shape) == 3:
        mask_np = mask_np[..., -1] # قناة ألفا
    ys, xs = np.where(mask_np > 0)
    if ys.size == 0 or xs.size == 0:
        return None
    return int(xs.min()), int(ys.min()), int(xs.max()), int(ys.max())

def crop_masked_object(img: Image.Image, mask: Image.Image):
    """اقطع المنطقة المحددة من الصورة حسب القناع وأرجعها"""
    bbox = get_mask_bbox(mask)
    if bbox is None:
        return None, None
    xmin, ymin, xmax, ymax = bbox
    # قص الصورة والقناع للمنطقة المحددة فقط
    cropped_img = img.crop((xmin, ymin, xmax + 1, ymax + 1))
    cropped_mask = mask.crop((xmin, ymin, xmax + 1, ymax + 1))
    # طبّق القناع على الجزء المقصوص (شفافية)
    cropped_img.putalpha(cropped_mask.split()[-1])
    return cropped_img, bbox

def paste_object_on_mask(
    base_img: Image.Image,
    base_mask: Image.Image,
    obj_img: Image.Image,
    obj_bbox,
):
    """الصق الجزء في منطقة القناع على الصورة الأساسية مع التحجيم والمحاذاة"""
    base_bbox = get_mask_bbox(base_mask)
    if base_bbox is None or obj_bbox is None:
        return None
    bxmin, bymin, bxmax, bymax = base_bbox
    ow, oh = obj_img.size
    # حجم منطقة اللصق في الصورة الأصلية
    bw, bh = bxmax - bxmin + 1, bymax - bymin + 1
    # غيّر حجم الجزء المقصوص ليناسب منطقة اللصق
    obj_img_resized = obj_img.resize((bw, bh), resample=Image.LANCZOS)
    # لصق الجزء على الصورة الأصلية في مكان القناع
    result = base_img.convert("RGBA")
    mask = base_mask.split()[-1]
    empty = Image.new("RGBA", result.size, (0, 0, 0, 0))
    empty.paste(obj_img_resized, (bxmin, bymin), obj_img_resized)
    # الدمج: الجزء الجديد مكان القناع فقط
    result = Image.composite(empty, result, mask)
    return result

def process_transfer(
    image1_editor, # dict: {"background": Image, "layers": [Image,...]}
    image2_editor  # dict: {"background": Image, "layers": [Image,...]}
):
    # التحقق من المدخلات
    if not image1_editor or not image2_editor:
        return None, None, None, "**❌ يرجى رفع الصورتين وتحديد الأقنعة."
    if (image1_editor.get("background") is None or
        not image1_editor.get("layers") or
        image2_editor.get("background") is None or
        not image2_editor.get("layers")
    ):
        return None, None, None, "**❌ يجب رسم قناع على كلتا الصورتين."
    base_img = image1_editor["background"].convert("RGBA")
    base_mask = image1_editor["layers"][0].convert("RGBA")
    obj_img = image2_editor["background"].convert("RGBA")
    obj_mask = image2_editor["layers"][0].convert("RGBA")

    # قص الجزء المحدد من الصورة الثانية
    cropped_obj, obj_bbox = crop_masked_object(obj_img, obj_mask)
    if cropped_obj is None:
        return None, None, None, "**❌ لم يتم تحديد أي كائن في الصورة الثانية."
    # لصق الجزء في منطقة القناع بالصورة الأولى
    result_img = paste_object_on_mask(base_img, base_mask, cropped_obj, obj_bbox)
    if result_img is None:
        return None, None, None, "**❌ لم يتم تحديد منطقة مناسبة للصق في الصورة الأولى."

    return base_mask, obj_mask, result_img, "**✅ تم نقل الكائن بنجاح!"

def create_interface():
    with gr.Blocks(title="نقل كائن من صورة إلى أخرى مع القناع والمحاذاة") as demo:
        gr.Markdown("""
        # 🖼️ نقل كائن من صورة إلى أخرى مع القناع والمحاذاة
        
        1. ارفع الصورة الأولى (صورة الخلفية).
        2. ارفع الصورة الثانية (الصورة التي بها الكائن).
        3. ارسم قناعًا على الكائن في الصورة الثانية (صورة الاستبدال).
        4. ارسم قناعًا في مكان اللصق في الصورة الأولى.
        5. اضغط زر "نقل الكائن" وستظهر النتيجة!
        """)
        with gr.Row():
            with gr.Column():
                image1_editor = gr.ImageMask(
                    label="الصورة الأولى (اختر مكان اللصق)",
                    type="pil",
                    height=400
                )
            with gr.Column():
                image2_editor = gr.ImageMask(
                    label="الصورة الثانية (حدد الكائن المراد نقله)",
                    type="pil",
                    height=400
                )
        with gr.Row():
            transfer_btn = gr.Button("🚀 نقل الكائن", variant="primary")
        with gr.Row():
            base_mask = gr.Image(label="قناع الصورة الأولى", type="pil", height=200)
            obj_mask = gr.Image(label="قناع الصورة الثانية", type="pil", height=200)
            result_img = gr.Image(label="النتيجة النهائية", type="pil", height=400)
        status = gr.Markdown()
        transfer_btn.click(
            fn=process_transfer,
            inputs=[image1_editor, image2_editor],
            outputs=[base_mask, obj_mask, result_img, status]
        )
    return demo

def main():
    demo = create_interface()
    demo.launch()

if __name__ == "__main__":
    main()