# -*- coding: utf-8 -*-
"""Face restoration.

This is the GFPGANer pipeline — detect, align, restore, paste back — rebuilt on
facexlib's FaceRestoreHelper plus a spandrel-loaded GFPGAN checkpoint, so the
abandoned gfpgan and basicsr packages are no longer needed.
"""

from __future__ import annotations

import threading

import numpy as np

from upscale import WEIGHTS_DIR, load_face_model, run_model

_helper = None
_helper_lock = threading.Lock()


def get_helper(upscale: int):
    """Returns a process-wide FaceRestoreHelper, retuned to the given upscale.

    The helper carries per-image state (landmarks, affines, cropped faces), so
    callers must hold `helper_lock()` for the whole detect->paste sequence.
    """
    global _helper
    import torch
    from facexlib.utils.face_restoration_helper import FaceRestoreHelper

    # facexlib's detection and parsing nets are kept on CUDA or CPU only. MPS is
    # deliberately excluded — these nets are untested there and the detector is
    # cheap relative to the SR pass, which still runs on the accelerator.
    det_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    if _helper is None:
        _helper = FaceRestoreHelper(
            upscale_factor=upscale,
            face_size=512,
            crop_ratio=(1, 1),
            det_model="retinaface_resnet50",
            save_ext="png",
            use_parse=True,
            device=det_device,
            model_rootpath=WEIGHTS_DIR,
        )
    else:
        _helper.set_upscale_factor(upscale)
    return _helper


def helper_lock() -> threading.Lock:
    return _helper_lock


def restore_faces(bgr: np.ndarray, background: np.ndarray, upscale: int) -> np.ndarray:
    """Restores every detected face in `bgr` and pastes them onto `background`.

    `bgr` is the original BGR image, `background` the already-upscaled BGR image
    the faces are composited onto. Returns BGR. If no face is found, the
    background is returned untouched.
    """
    model = load_face_model()

    with _helper_lock:
        helper = get_helper(upscale)
        helper.clean_all()
        helper.read_image(bgr)
        helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
        helper.align_warp_face()

        if not helper.cropped_faces:
            return background

        for cropped in helper.cropped_faces:
            # cropped is 512x512 BGR uint8; the model is scale 1.
            rgb = cropped[:, :, ::-1]
            restored = run_model(model, np.ascontiguousarray(rgb), tile=0)
            helper.add_restored_face(np.ascontiguousarray(restored[:, :, ::-1]))

        helper.get_inverse_affine(None)
        return helper.paste_faces_to_input_image(upsample_img=background)