from base64 import b64encode
from io import BytesIO
from pathlib import Path

import numpy as np
from basicsr.archs.rrdbnet_arch import RRDBNet
from PIL import Image
from realesrgan import RealESRGANer


class EndpointHandler:
    def __init__(self, path=""):
        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
        self.upsampler = RealESRGANer(
            scale=4,
            model_path=str(Path(path) / "RealESRGAN_x4plus.pth"),
            model=model,
            tile=0,
            tile_pad=10,
            pre_pad=0,
            half=True,
        )

    def __call__(self, data):
        """
        Args:
            data (:obj:):
                includes the input data and the parameters for the inference.
        Return:
            A :obj:`dict`:. base64 encoded image
        """
        image = data.pop("inputs", data)
        # image = Image.open(BytesIO(image)).convert("RGB")
        image = np.array(image)
        image = image[:, :, ::-1]  # RGB -> BGR

        image, _ = self.upsampler.enhance(image, outscale=4)
        image = image[:, :, ::-1]  # BGR -> RGB
        image = Image.fromarray(image)

        # encode image as base 64
        buffered = BytesIO()
        image.save(buffered, format="PNG")
        img_str = b64encode(buffered.getvalue())

        # postprocess the prediction
        return {"image": img_str.decode()}