#!/usr/bin/env python3

import os
import logging

import cv2
import random  # TODO: seed!

import insightface
from insightface.app import FaceAnalysis

from insightface.model_zoo.inswapper import INSwapper
from insightface.app.common import Face

logger = logging.getLogger("protestswap")

SWAPPING_MODEL = "inswapper_128.onnx"


class ProtestFaceSwapper:
    def __init__(self, root_dir):
        self._root_dir = root_dir
        self._app: FaceAnalysis = FaceAnalysis(name="buffalo_l", root=root_dir)
        self._model: INSwapper
        self._target: cv2.typing.MatLike | None = None
        self._target_faces: list[Face] = []
        self._template_faces: list[Face] = []

        self._app.prepare(ctx_id=0, det_size=(640, 640))

    def prepare_model(self) -> None:
        self._model = insightface.model_zoo.get_model(
            os.path.join(self._root_dir, "models", SWAPPING_MODEL),
            root=self._root_dir,
            download=False,
            download_zip=False,
        )  # Revealed type: "INSwapper"
        print(f"Model: {self._model}")

    def set_target_path(self, path: str) -> None:
        self.set_target(cv2.imread(path))

    def set_target(self, image: cv2.UMat) -> None:
        self._target = image

    def detect_target_faces(self):
        self._target_faces = self._app.get(self._target)
        print(f"Found {len(self._target_faces)} faces!")

    def add_template_path(self, template_path: str):
        self.add_template(cv2.imread(template_path))

    def add_template(self, template: cv2.UMat):
        self._template_faces += self._app.get(template)

    def swap(self) -> cv2.typing.MatLike:
        work_copy = self._target.copy()
        for i, face in enumerate(self._target_faces):
            print(
                f"Replacing face {i + 1}/{len(self._target_faces)}...",
                # end="",
                # flush=True,
            )
            imposed_face = random.choice(self._template_faces)
            print(
                f"work_copy: {type(work_copy)}, face: {type(face)}, imposed: {
                    type(imposed_face)
                }"
            )
            work_copy = self._model.get(
                work_copy, face, imposed_face, paste_back=True)
            print(" done!")

        return work_copy

    @staticmethod
    def readFile(path: str) -> cv2.typing.MatLike | None:
        return cv2.imread(path)
