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


class FaceSwapper:
    def __init__(self, det_size=(640, 640), ctx_id=0):
        """
        Initialize face analysis + inswapper on GPU when available.
        ctx_id=0 prefers CUDA. Falls back gracefully if no GPU.
        """
        available = ort.get_available_providers()
        providers = []
        if "CUDAExecutionProvider" in available:
            providers.append("CUDAExecutionProvider")
        providers.append("CPUExecutionProvider")

        self.app = FaceAnalysis(
            name="buffalo_l",
            providers=providers
        )
        self.app.prepare(ctx_id=ctx_id, det_size=det_size)

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

        print(f"[FaceSwapper] Providers in use: {providers}")
        print(f"[FaceSwapper] Available ORT providers: {available}")

    def swap_faces(self, source_path, source_face_idx, target_path, target_face_idx):
        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-to-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, "
                f"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, "
                f"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
        )
        return result

    def count_faces(self, img_path):
        """Count faces using the same InsightFace detector (GPU-aware)."""
        img = cv2.imread(img_path)
        if img is None:
            return 0
        faces = self.app.get(img)
        return len(faces)