import base64
import time
import sys

import numpy as np
import triton_python_backend_utils as pb_utils  # type: ignore
if "/" not in sys.path:
    sys.path.insert(0, "/")
from utils import decode_image, od_resize_image


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

        responses = []
        st = time.time()
        for request in requests:
            image_bytes = pb_utils.get_input_tensor_by_name(request, "image_bytes").as_numpy()
            image_bytes = base64.b64decode(image_bytes[0].decode('utf-8'))
            image = decode_image(image_bytes)

            (
                resized_image,
                (original_width, original_height),
                (resized_width, resized_height),
            ) = od_resize_image(image, (960, 960))

            resized_image = resized_image[:, :, ::-1].transpose(2, 0, 1)
            resized_image = np.ascontiguousarray(resized_image).astype(np.float32) / 255.0
            resized_image = np.expand_dims(resized_image, axis=0)

            out_tensor_0 = pb_utils.Tensor("resized_image", np.array(resized_image))
            out_tensor_1 = pb_utils.Tensor("original_shape", np.array([original_width, original_height]).astype(np.uint16))
            out_tensor_2 = pb_utils.Tensor("resized_shape", np.array([960, 960]).astype(np.uint16))

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

        logger.log_info(f"OD Preprocess execute duration : {int((time.time() - st)*1000)} ms")
        return responses