from typing import Dict, List, Any
from gfpgan import GFPGANer
import cv2
from imageio import imread
from basicsr.utils import imwrite
import io
import os
import numpy as np
import base64


class EndpointHandler():
    def __init__(self, path=""):
        self.restorer = GFPGANer(
        model_path="./GFPGANv1.4.pth",
        upscale=2,
        arch="clean",
        channel_multiplier=2,
        bg_upsampler=None)
    
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
       data args:
            inputs (:obj: `str`)
            date (:obj: `str`)
      Return:
            A :obj:`list` | `dict`: will be serialized and returned
        """
        # get inputs
        inputs = data.pop("inputs",data)
        img = imread(io.BytesIO(base64.b64decode(inputs)))
        cropped_faces, restored_faces, restored_img = self.restorer.enhance(
            img,
            has_aligned=False,
            only_center_face=False,
            paste_back=True,
            weight=0.5)

        for idx, (cropped_face, restored_face) in enumerate(zip(cropped_faces, restored_faces)):
            retval, buffer = cv2.imencode('.png', restored_face)
            jpg_as_text = base64.b64encode(buffer)
            return jpg_as_text