import cv2
import insightface
import numpy as np
import os
from gfpgan import GFPGANer
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer   # pip install realesrgan


class FaceSwapper:
    def __init__(self,
                 model_path="models/inswapper_128.onnx",
                 gfpgan_model_path="gfpgan/weights/GFPGANv1.4.pth",
                 realesrgan_model_path="models/RealESRGAN_x2plus.pth"):
        
        # ============ Face Analysis ============
        self.app = insightface.app.FaceAnalysis(name="buffalo_l", providers=['CPUExecutionProvider'])
        self.app.prepare(ctx_id=0, det_size=(640, 640))  # حافظ عليها صغيرة

        # ============ Face Swapper ============
        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=2,   # ممكن تخليها 4 حسب جهازك
            arch="clean",
            channel_multiplier=2
        )

        # ============ Real-ESRGAN ============
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                        num_block=23, num_grow_ch=32, scale=2)
        self.realesrganer = RealESRGANer(
            scale=2,
            model_path=realesrgan_model_path,
            model=model,
            tile=0,
            tile_pad=10,
            pre_pad=0,
            half=False
        )

    @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:
            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, (51, 51), 40)
        center = ((x1 + x2) // 2, (y1 + y2) // 2)

        blended = cv2.seamlessClone(face_crop, swapped_img, mask, center, cv2.MIXED_CLONE)

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

        # تكبير وتحسين الصورة كاملة بالـ Real-ESRGAN
        sr_img, _ = self.realesrganer.enhance(gfpgan_img, outscale=1)

        cv2.imwrite(output_path, sr_img, [cv2.IMWRITE_PNG_COMPRESSION, 0])
        return output_path