from chain_img_processor import ChainImgProcessor, ChainImgPlugin
import os
import gfpgan
import threading
from PIL import Image
from numpy import asarray
import cv2

from roop.utilities import resolve_relative_path, conditional_download
modname = os.path.basename(__file__)[:-3] # calculating modname

model_gfpgan = None
THREAD_LOCK_GFPGAN = threading.Lock()


# start function
def start(core:ChainImgProcessor):
    manifest = { # plugin settings
        "name": "GFPGAN", # name
        "version": "1.4", # version

        "default_options": {},
        "img_processor": {
            "gfpgan": GFPGAN
        }
    }
    return manifest

def start_with_options(core:ChainImgProcessor, manifest:dict):
    pass


class GFPGAN(ChainImgPlugin):

    def init_plugin(self):
        global model_gfpgan

        if model_gfpgan is None:
            model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
            model_gfpgan = gfpgan.GFPGANer(model_path=model_path, upscale=1, device=self.device) # type: ignore[attr-defined]



    def process(self, frame, params:dict):
        import copy

        global model_gfpgan

        if model_gfpgan is None:
            return frame 
        
        if "face_detected" in params:
            if not params["face_detected"]:
                return frame
        # don't touch original    
        temp_frame = copy.copy(frame)
        if "processed_faces" in params:
            for face in params["processed_faces"]:
                start_x, start_y, end_x, end_y = map(int, face['bbox'])
                padding_x = int((end_x - start_x) * 0.5)
                padding_y = int((end_y - start_y) * 0.5)
                start_x = max(0, start_x - padding_x)
                start_y = max(0, start_y - padding_y)
                end_x = max(0, end_x + padding_x)
                end_y = max(0, end_y + padding_y)
                temp_face = temp_frame[start_y:end_y, start_x:end_x]
                if temp_face.size:
                    with THREAD_LOCK_GFPGAN:
                        _, _, temp_face = model_gfpgan.enhance(
                                temp_face,
                                paste_back=True
                            )
                    temp_frame[start_y:end_y, start_x:end_x] = temp_face
        else:
            with THREAD_LOCK_GFPGAN:
                _, _, temp_frame = model_gfpgan.enhance(
                        temp_frame,
                        paste_back=True
                    )

        if not "blend_ratio" in params: 
            return temp_frame

        temp_frame = Image.blend(Image.fromarray(frame), Image.fromarray(temp_frame), params["blend_ratio"])
        return asarray(temp_frame)