import gradio as gr
from ultralytics import YOLO
import cv2
import numpy as np
from PIL import Image

print("🚀 Loading Manga Sniper Studio (Optimized for Auto-Cleaning)...")
# تحميل نموذج YOLO الخاص بكشف فقاعات المانجا
model = YOLO("best.pt")

def clean_manga_page(input_image):
    try:
        if input_image is None:
            return None
        
        # تحويل الصورة المدخلة إلى مصفوفة numpy بنظام BGR لـ OpenCV
        img_cv = cv2.cvtColor(np.array(input_image), cv2.COLOR_RGB2BGR)
        h, w = img_cv.shape[:2]
        
        # 1. كشف مناطق الفقاعات باستخدام YOLO
        results = model(img_cv, conf=0.35)
        bubble_mask = np.zeros((h, w), dtype=np.uint8)

        if results[0].masks is not None:
            for seg in results[0].masks.xy:
                poly = seg.astype(np.int32)
                # فلتر حجم الفقاعة لتفادي كشف عناصر الصفحة بالخطأ
                if 300 < cv2.contourArea(poly) < (h * w * 0.50):
                    cv2.fillPoly(bubble_mask, [poly], 255)
        
        # 2. استخراج الحبر والنصوص داخل الفقاعات بدقة عالية
        gray = cv2.cvtColor(img_cv, cv2.COLOR_BGR2GRAY)
        
        # استخدام عتبة تكيفية متطورة لالتقاط النصوص بوضوح
        ink = cv2.adaptiveThreshold(
            gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
            cv2.THRESH_BINARY_INV, 15, 4
        )
        
        # حصر الحبر المستخرج داخل قناع الفقاعات فقط
        ink_in_bubble = cv2.bitwise_and(ink, bubble_mask)

        # 3. توسيع ذكي للحواف (Dilation) لضمان إخفاء النص بالكامل وعدم ترك هوامش سوداء
        kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (4, 4))
        final_mask = cv2.dilate(ink_in_bubble, kernel, iterations=2)

        # 4. عملية التبييض وإعادة البناء الاحترافية (Inpainting)
        # نستخدم خوارزمية Telea السريعة والدقيقة لمسح النصوص وتعويضها بالخلفية البيضاء للفقاعة
        cleaned_cv = cv2.inpaint(img_cv, final_mask, inpaintRadius=5, flags=cv2.INPAINT_TELEA)
        
        # تحويل الصورة الناتجة مجدداً إلى نظام RGB لتعرض في Gradio بدون خربطة ألوان
        cleaned_rgb = cv2.cvtColor(cleaned_cv, cv2.COLOR_BGR2RGB)
        return Image.fromarray(cleaned_rgb)
        
    except Exception as e:
        print(f"Error occurred: {str(e)}")
        return input_image

# تصميم واجهة مستخدم احترافية وشبيهة بالتطبيقات المتطورة
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🌌 AI Manga Cleaner Studio
        قم برفع صفحة المانجا/المانهوا الخاصة بك ليقوم الذكاء الاصطناعي بتبييض ومسح النصوص تلقائياً مع الحفاظ على الرسوم الخلفية.
        """
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil", label="الصورة الأصلية (Original Page)")
            clean_btn = gr.Button("✨ بدء التبييض التلقائي", variant="primary")
        with gr.Column():
            image_output = gr.Image(type="pil", label="النتيجة المبيّضة (Cleaned Page)")
            
    clean_btn.click(fn=clean_manga_page, inputs=image_input, outputs=image_output)

# تشغيل التطبيق
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)