from __future__ import annotations

import logging
import os
from functools import cached_property
from typing import TYPE_CHECKING, Callable

import cv2
import numpy as np
import torch

from modules import devices, errors, face_restoration, shared

if TYPE_CHECKING:
    from facexlib.utils.face_restoration_helper import FaceRestoreHelper

logger = logging.getLogger(__name__)


def create_face_helper(device) -> FaceRestoreHelper:
    from facexlib.detection import retinaface
    from facexlib.utils.face_restoration_helper import FaceRestoreHelper
    if hasattr(retinaface, 'device'):
        retinaface.device = device
    return FaceRestoreHelper(
        upscale_factor=1,
        face_size=512,
        crop_ratio=(1, 1),
        det_model='retinaface_resnet50',
        save_ext='png',
        use_parse=True,
        device=device,
    )


def restore_with_face_helper(
    np_image: np.ndarray,
    face_helper: FaceRestoreHelper,
    restore_face: Callable[[torch.Tensor], torch.Tensor],
) -> np.ndarray:
    """
    Find faces in the image using face_helper, restore them using restore_face, and paste them back into the image.

    `restore_face` should take a cropped face image and return a restored face image.
    """
    from basicsr.utils import img2tensor, tensor2img
    from torchvision.transforms.functional import normalize
    np_image = np_image[:, :, ::-1]
    original_resolution = np_image.shape[0:2]

    try:
        logger.debug("Detecting faces...")
        face_helper.clean_all()
        face_helper.read_image(np_image)
        face_helper.get_face_landmarks_5(only_center_face=False, resize=640, eye_dist_threshold=5)
        face_helper.align_warp_face()
        logger.debug("Found %d faces, restoring", len(face_helper.cropped_faces))
        for cropped_face in face_helper.cropped_faces:
            cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
            normalize(cropped_face_t, (0.5, 0.5, 0.5), (0.5, 0.5, 0.5), inplace=True)
            cropped_face_t = cropped_face_t.unsqueeze(0).to(devices.device_codeformer)

            try:
                with torch.no_grad():
                    restored_face = tensor2img(
                        restore_face(cropped_face_t),
                        rgb2bgr=True,
                        min_max=(-1, 1),
                    )
                devices.torch_gc()
            except Exception:
                errors.report('Failed face-restoration inference', exc_info=True)
                restored_face = tensor2img(cropped_face_t, rgb2bgr=True, min_max=(-1, 1))

            restored_face = restored_face.astype('uint8')
            face_helper.add_restored_face(restored_face)

        logger.debug("Merging restored faces into image")
        face_helper.get_inverse_affine(None)
        img = face_helper.paste_faces_to_input_image()
        img = img[:, :, ::-1]
        if original_resolution != img.shape[0:2]:
            img = cv2.resize(
                img,
                (0, 0),
                fx=original_resolution[1] / img.shape[1],
                fy=original_resolution[0] / img.shape[0],
                interpolation=cv2.INTER_LINEAR,
            )
        logger.debug("Face restoration complete")
    finally:
        face_helper.clean_all()
    return img


class CommonFaceRestoration(face_restoration.FaceRestoration):
    net: torch.Module | None
    model_url: str
    model_download_name: str

    def __init__(self, model_path: str):
        super().__init__()
        self.net = None
        self.model_path = model_path
        os.makedirs(model_path, exist_ok=True)

    @cached_property
    def face_helper(self) -> FaceRestoreHelper:
        return create_face_helper(self.get_device())

    def send_model_to(self, device):
        if self.net:
            logger.debug("Sending %s to %s", self.net, device)
            self.net.to(device)
        if self.face_helper:
            logger.debug("Sending face helper to %s", device)
            self.face_helper.face_det.to(device)
            self.face_helper.face_parse.to(device)

    def get_device(self):
        raise NotImplementedError("get_device must be implemented by subclasses")

    def load_net(self) -> torch.Module:
        raise NotImplementedError("load_net must be implemented by subclasses")

    def restore_with_helper(
        self,
        np_image: np.ndarray,
        restore_face: Callable[[torch.Tensor], torch.Tensor],
    ) -> np.ndarray:
        try:
            if self.net is None:
                self.net = self.load_net()
        except Exception:
            logger.warning("Unable to load face-restoration model", exc_info=True)
            return np_image

        try:
            self.send_model_to(self.get_device())
            return restore_with_face_helper(np_image, self.face_helper, restore_face)
        finally:
            if shared.opts.face_restoration_unload:
                self.send_model_to(devices.cpu)


def patch_facexlib(dirname: str) -> None:
    import facexlib.detection
    import facexlib.parsing

    det_facex_load_file_from_url = facexlib.detection.load_file_from_url
    par_facex_load_file_from_url = facexlib.parsing.load_file_from_url

    def update_kwargs(kwargs):
        return dict(kwargs, save_dir=dirname, model_dir=None)

    def facex_load_file_from_url(**kwargs):
        return det_facex_load_file_from_url(**update_kwargs(kwargs))

    def facex_load_file_from_url2(**kwargs):
        return par_facex_load_file_from_url(**update_kwargs(kwargs))

    facexlib.detection.load_file_from_url = facex_load_file_from_url
    facexlib.parsing.load_file_from_url = facex_load_file_from_url2