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


class FaceSwapper:
    def __init__(self, model_name: str = "buffalo_l", det_size=(640, 640)):
        self.providers = self._resolve_providers()
        self.app = FaceAnalysis(name=model_name, providers=self.providers)
        ctx_id = 0 if "CUDAExecutionProvider" in self.providers else -1
        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=self.providers,
        )

    @staticmethod
    def _resolve_providers():
        available = ort.get_available_providers()
        if "CUDAExecutionProvider" in available:
            return ["CUDAExecutionProvider", "CPUExecutionProvider"]
        return ["CPUExecutionProvider"]

    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 = sorted(self.app.get(source_img), key=lambda x: x.bbox[0])
        target_faces = sorted(self.app.get(target_img), 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]
        return self.swapper.get(target_img, target_face, source_face, paste_back=True)

    def count_faces(self, img_path):
        img = cv2.imread(img_path)
        if img is None:
            raise ValueError(f"Could not read image: {img_path}")
        return len(self.app.get(img))