import cv2
import insightface
from insightface.app import FaceAnalysis
import onnxruntime as ort
import os
import numpy as np

class FaceSwapper:
    def __init__(self, det_size=(1280, 1280), use_enhancer=False):
        """
        det_size: higher = better detection on A100 (try 1280 or even 1600)
        use_enhancer: enable GFPGAN-style enhancement (requires gfpgan package)
        """
        # Force CUDA first
        available = ort.get_available_providers()
        print(f"[FaceSwapper] Available ONNX providers: {available}")

        providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']
        if 'CUDAExecutionProvider' not in available:
            print("[WARNING] CUDAExecutionProvider not found! Falling back to CPU. Check onnxruntime-gpu + CUDA.")
            providers = ['CPUExecutionProvider']

        # Face analysis (detection + landmarks + embedding)
        self.app = FaceAnalysis(
            name='buffalo_l',
            providers=providers
        )
        # ctx_id=0 → first GPU. A100 loves larger det_size
        self.app.prepare(ctx_id=0, det_size=det_size)

        # Inswapper
        self.swapper = insightface.model_zoo.get_model(
            'inswapper_128.onnx',
            download=True,
            download_zip=True,
            providers=providers
        )

        self.use_enhancer = false
        self.enhancer = None
        if use_enhancer:
            try:
                from gfpgan import GFPGANer
                # You need to place GFPGANv1.4.pth in the Space or let it download
                self.enhancer = GFPGANer(
                    model_path='https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth',
                    upscale=1,
                    arch='clean',
                    channel_multiplier=2,
                    bg_upsampler=None
                )
                print("[FaceSwapper] GFPGAN enhancer loaded")
            except Exception as e:
                print(f"[WARNING] Could not load enhancer: {e}")
                self.use_enhancer = False

    def swap_faces(self, source_path, source_face_idx, target_path, target_face_idx, enhance=False):
        source_img = cv2.imread(source_path)
        target_img = cv2.imread(target_path)

        if source_img is None or target_img is None:
            raise ValueError("Could not read one or both images")

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

        # Sort left → right for consistent indexing
        source_faces = sorted(source_faces, key=lambda x: x.bbox[0])
        target_faces = sorted(target_faces, key=lambda x: x.bbox[0])

        if len(source_faces) < source_face_idx or source_face_idx < 1:
            raise ValueError(f"Source image contains {len(source_faces)} faces, but requested face {source_face_idx}")
        if len(target_faces) < target_face_idx or target_face_idx < 1:
            raise ValueError(f"Target image contains {len(target_faces)} faces, but requested face {target_face_idx}")

        source_face = source_faces[source_face_idx - 1]
        target_face = target_faces[target_face_idx - 1]

        result = self.swapper.get(target_img, target_face, source_face, paste_back=True)

        # Optional enhancement (big quality jump)
        if (enhance or self.use_enhancer) and self.enhancer is not None:
            _, _, result = self.enhancer.enhance(
                result,
                has_aligned=False,
                only_center_face=False,
                paste_back=True,
                weight=0.8   # 0.5–1.0, higher = stronger restoration
            )

        return result

    def count_faces(self, img_path):
        img = cv2.imread(img_path)
        if img is None:
            return 0
        faces = self.app.get(img)
        return len(faces)