import base64

import cv2
import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore


class TritonPythonModel:
    def initialize(self, args):
        self.td_canvas_min_size = 640
        self.td_canvas_max_size = 2240

    def _preprocess(self, image_bytes: bytes):
        image_array = np.frombuffer(image_bytes, dtype=np.uint8)
        original_image = cv2.imdecode(image_array, cv2.IMREAD_COLOR_BGR)
        if original_image.shape[2] == 4:
            original_image = original_image[:, :, :3]
        h, w = original_image.shape[:2]

        if min(h, w) < 640:
            max_ratio = 1.0
        else:
            max_ratio = 1.5

        ratio = min(self.td_canvas_max_size / max(h,w), max_ratio)

        if ratio != 1.0:
            ratio = float(int(ratio * 32)) / 32.0
            resized_image = cv2.resize(original_image, None, fx=ratio, fy=ratio, interpolation=cv2.INTER_CUBIC)
        else:
            resized_image = original_image

        h, w = resized_image.shape[:2]
        if h < self.td_canvas_min_size:
            h_pad = self.td_canvas_min_size - h
        else:
            h_pad = min(64 - h%64, self.td_canvas_max_size - h)

        if w < self.td_canvas_min_size:
            w_pad = self.td_canvas_min_size - w
        else:
            w_pad = min(64 - w%64, self.td_canvas_max_size - w)
        resized_image = np.pad(resized_image, ((0, (h_pad)), (0, w_pad), (0, 0)), 'constant', constant_values=0)

        resized_image = resized_image.transpose(2, 0, 1)
        resized_image = np.expand_dims(resized_image, axis=0)
        resized_image = np.ascontiguousarray(resized_image)

        out_tensor_0 = pb_utils.Tensor("resized_image", resized_image)
        out_tensor_1 = pb_utils.Tensor("resize_ratio", np.array([ratio], dtype=np.float32))
        out_tensor_2 = pb_utils.Tensor("original_image", original_image)

        return out_tensor_0, out_tensor_1, out_tensor_2

    def execute(self, requests):
        logger = pb_utils.Logger

        responses = []

        for request in requests:
            try:
                input_tensor = pb_utils.get_input_tensor_by_name(request, "image_bytes").as_numpy()
                image_bytes = base64.b64decode(input_tensor[0].decode('utf-8'))
                out_tensor_0, out_tensor_1, out_tensor_2 = self._preprocess(image_bytes)

            except pb_utils.TritonModelException as e:
                responses.append(pb_utils.InferenceResponse(error=pb_utils.TritonError(str(e), pb_utils.TritonError.BAD_REQUEST)))
                continue

            responses.append(pb_utils.InferenceResponse(output_tensors=[out_tensor_0, out_tensor_1, out_tensor_2]))

        return responses