import os
import numpy as np
from typing import Dict

from transformers import MaskFormerFeatureExtractor
# from refactor.utils.assets import ImageAsset
from io import BytesIO
from PIL import Image
import rawpy

import triton_python_backend_utils as pb_utils


class TritonPythonModel:
    def initialize(self, args: Dict[str, str]) -> None:
        """
        Initialize the tokenization process
        :param args: arguments from Triton config file
        """
        path: str = os.path.join(args["model_repository"], args["model_version"])
        self.processor = MaskFormerFeatureExtractor.from_pretrained(path)
        
    def execute(self, requests):
        """
        Parameters
        ----------
        requests : list
          A list of pb_utils.InferenceRequest

        Returns
        -------
        list
          A list of pb_utils.InferenceResponse. The length of this list must
          be the same as `requests`
        """
        responses = []
        # loop batch request
        for request in requests:
            # Get INPUT
            # img = [
            #     self.preprocess(t)
            #     for t in pb_utils.get_input_tensor_by_name(request, "IMAGE")
            #     .as_numpy()
            # ]
            imgs = pb_utils.get_input_tensor_by_name(request, "IMAGE").as_numpy().tolist()
            imgs = self.preprocess(imgs)
            # imgs = [Image.open(BytesIO(img)) for img in imgs]
            # img = ImageAsset(inp).image
            # img = [self.preprocess(e) for e in inp]

            # Get EXTRACTFEATURE
            feature: Dict[str, np.ndarray] = self.processor(imgs, return_tensors='np')
            outputs = pb_utils.Tensor('pixel_values', feature['pixel_values'])
            inference_response = pb_utils.InferenceResponse(output_tensors=[outputs])
            responses.append(inference_response)
            
        return responses
    

    def preprocess(self, data):
      print("images: ", data)
      imgs = []
      for e in data:
        buffer = BytesIO(e)
        print("Type of data: ", type(buffer))
        _image = Image.open(buffer).convert('RGB')
        print("Type of _image: ", type(_image))
        imgs.append(_image)
      print("Load success")
      # try:
      #   with rawpy.imread(buffer) as raw:
      #     thumb = raw.extract_thumb()
      #     print("Load success: ",thumb)
      #   if thumb.format == rawpy.ThumbFormat.JPEG or thumb.format == rawpy.ThumbFormat.PNG:
      #     _image = Image.open(thumb.data)
      #     print("PNG Type of _image: ", type(_image))
      #   elif thumb.format == rawpy.ThumbFormat.BITMAP:
      #     _image = Image.fromarray(thumb.data)
      #     print("bitmap Type of _image: ", type(_image))
      # except:
      #   raise AssertionError(f'Cannot load image asset from {str(buffer):20s}')
      return imgs
      
    def finalize(self) -> None:
        """
        Finalize the model
        """
        pass