import time

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


class TritonPythonModel:
    @staticmethod
    def _to_bytes(value) -> bytes:
        if isinstance(value, np.bytes_):
            return value.tobytes()
        if isinstance(value, bytes):
            return value
        if isinstance(value, bytearray):
            return bytes(value)
        if isinstance(value, str):
            return value.encode("utf-8")
        return bytes(value)

    def execute(self, requests):
        logger = pb_utils.Logger
        responses = []
        st = time.time()

        for request in requests:
            image_bytes_tensor = pb_utils.get_input_tensor_by_name(request, "IMAGE_BYTES_LIST")
            image_bytes_list = image_bytes_tensor.as_numpy().reshape(-1)

            images = []
            for idx, image_bytes in enumerate(image_bytes_list):
                raw_bytes = self._to_bytes(image_bytes)
                if len(raw_bytes) == 0:
                    raise pb_utils.TritonModelException(f"Empty image bytes at index {idx}.")

                image_np = np.frombuffer(raw_bytes, np.uint8)
                image = cv2.imdecode(image_np, cv2.IMREAD_COLOR)
                if image is None:
                    raise pb_utils.TritonModelException(
                        f"Failed to decode image bytes at index {idx}."
                    )

                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                image = cv2.resize(image, (224, 224), interpolation=cv2.INTER_AREA)

                x = image.astype(np.float32) / 255.0
                x = (x - 0.5) / 0.5
                x = np.transpose(x, (2, 0, 1))
                images.append(x)

            if len(images) == 0:
                pixel_values = np.empty((0, 3, 224, 224), dtype=np.float32)
            else:
                pixel_values = np.stack(images, axis=0).astype(np.float32)

            responses.append(
                pb_utils.InferenceResponse(
                    output_tensors=[pb_utils.Tensor("pixel_values", pixel_values)]
                )
            )

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