import gradio as gr
import os
import uuid
import cv2
from mask import FaceSwapper  # الكلاس اللي عملناه قبل كده

# تحميل الموديل مره واحده
swapper = FaceSwapper(
    model_path="models/inswapper_128.onnx",
    gfpgan_model_path="gfpgan/weights/GFPGANv1.4.pth"
)

# دالة للتعامل مع Gradio
def swap_faces(source_img, target_img):
    try:
        # حفظ الصور المرفوعة مؤقتاً
        source_path = f"temp_source_{uuid.uuid4().hex}.jpg"
        target_path = f"temp_target_{uuid.uuid4().hex}.jpg"
        output_path = f"img/result_{uuid.uuid4().hex}.jpg"

        cv2.imwrite(source_path, cv2.cvtColor(source_img, cv2.COLOR_RGB2BGR))
        cv2.imwrite(target_path, cv2.cvtColor(target_img, cv2.COLOR_RGB2BGR))

        # تنفيذ face swap
        result_path = swapper.merge_face_into_image(source_path, target_path, output_path)

        # تحميل الصورة الناتجة وعرضها
        result_img = cv2.cvtColor(cv2.imread(result_path), cv2.COLOR_BGR2RGB)

        # مسح الملفات المؤقتة
        os.remove(source_path)
        os.remove(target_path)

        return result_img
    except Exception as e:
        return f"❌ حصل خطأ: {str(e)}"

# واجهة Gradio
demo = gr.Interface(
    fn=swap_faces,
    inputs=[
        gr.Image(type="numpy", label="Source Image (الطفل)"),
        gr.Image(type="numpy", label="Target Image (المشهد)")
    ],
    outputs=gr.Image(type="numpy", label="النتيجة"),
    title="FaceSwap with GFPGAN",
    description="ارفع صورتين: (1) صورة الطفل (2) المشهد اللي عايز تدخله فيه. وهنرجعلك صورة face swap محسّنة بـ GFPGAN."
)

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


"""import cv2
import insightface
import numpy as np
import os
from gfpgan import GFPGANer  # pip install gfpgan"""


"""class FaceSwapper:
    def __init__(self,
                 model_path="models/inswapper_128.onnx",
                 gfpgan_model_path="gfpgan/weights/GFPGANv1.4.pth"):
        # ============ تحميل FaceAnalysis ============
        self.app = insightface.app.FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))

        # ============ تحميل inswapper ============
        if not os.path.exists(model_path):
            raise FileNotFoundError(f"❌ الموديل مش موجود في: {model_path}")
        self.swapper = insightface.model_zoo.get_model(model_path, providers=['CPUExecutionProvider'])

        # ============ تحميل GFPGAN ============
        self.gfpganer = GFPGANer(
            model_path=gfpgan_model_path,
            upscale=1,
            arch="clean",
            channel_multiplier=2
        )

    # ============ دالة مساعدة لاختيار أكبر وجه ============
    @staticmethod
    def get_biggest_face(faces):
        return max(faces, key=lambda f: (f.bbox[2]-f.bbox[0]) * (f.bbox[3]-f.bbox[1]))

    # ============ دالة الدمج ============
    def merge_face_into_image(self, source_img_path, target_img_path, output_path):
        source_img = cv2.imread(source_img_path)
        target_img = cv2.imread(target_img_path)

        if source_img is None or target_img is None:
            raise ValueError("❌ مشكلة في قراءة الصور")

        source_faces = self.app.get(source_img)
        target_faces = self.app.get(target_img)

        if not source_faces or not target_faces:
            #raise ValueError("❌ مش لاقي وش في الصورة!")
            print("⚠️ No faces detected, returning target image as-is.")
            cv2.imwrite(output_path, target_img)
            return output_path

        source_face = self.get_biggest_face(source_faces)
        target_face = self.get_biggest_face(target_faces)

        # استبدال الوجه
        swapped_img = self.swapper.get(target_img.copy(), target_face, source_face, paste_back=True)

        # قص الوجه من الصورة
        x1, y1, x2, y2 = target_face.bbox.astype(int)
        x1, y1 = max(0, x1), max(0, y1)
        x2, y2 = min(swapped_img.shape[1], x2), min(swapped_img.shape[0], y2)

        face_crop = swapped_img[y1:y2, x1:x2]

        if face_crop.size == 0:
            raise ValueError("❌ الوجه المقطوع فاضي (bbox مش مظبوط)")

        # ماسك بنفس حجم الوجه
        mask = 255 * np.ones(face_crop.shape, face_crop.dtype)
        mask = cv2.GaussianBlur(mask, (25, 25), 30)

        # مركز الوجه
        center = ((x1 + x2) // 2, (y1 + y2) // 2)

        # دمج الوجه في الصورة
        blended = cv2.seamlessClone(face_crop, swapped_img, mask, center, cv2.NORMAL_CLONE)

        # تحسين الصورة بالـ GFPGAN
        _, _, enhanced = self.gfpganer.enhance(blended, has_aligned=False, only_center_face=False, paste_back=True)

        final_img = enhanced  # أو blended لو GFPGAN مش موجود
        return cv2.cvtColor(final_img, cv2.COLOR_BGR2RGB)"""