# model.py
import triton_python_backend_utils as pb_utils
import numpy as np
import json
import random
import string

class TritonPythonModel:
    def initialize(self, args):
        """
        모델이 로드될 때 딱 한 번만 호출됩니다.
        `initialize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
        이 모델과 관련된 모든 상태를 초기화할 수 있습니다.
        """
        self.logger = pb_utils.Logger

        self.model_name = args["model_name"]
        self.model_config = json.loads(args["model_config"])

        self.logger.log_info(f"'{self.model_name}' 모델 초기화 완료")


    def execute(self, requests):
        """
        Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다.
        """
        responses = []

        # 각 추론 요청을 순회하며 처리합니다.
        for request in requests:
            # Triton 입력 파싱            
            input_text = self._get_input_value(request, "INPUT")
            self.logger.log_info(f"INPUT 출력:\n{input_text}")
            
            random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))

            output = random_string
            self.logger.log_info(f"OUTPUT 출력:\n{output}")

            # 생성된 텍스트를 Triton 출력 텐서로 변환합니다.
            output_tensor = pb_utils.Tensor("OUTPUT", np.array(output.encode('utf-8'), dtype=np.bytes_))

            # 응답 객체를 생성하고 출력 텐서를 추가합니다.
            responses.append(pb_utils.InferenceResponse(output_tensors=[output_tensor]))

        return responses

    # 수정 후
    def _get_input_value(self, request, input_name: str, default=None):
        tensor_value = pb_utils.get_input_tensor_by_name(request, input_name)
        if tensor_value is None:
            return default
        
        # 텐서의 데이터 타입을 확인
        # bytes 타입의 경우, as_numpy()를 호출하면 NumPy 배열이 아닌 bytearray가 반환될 수 있음
        # 따라서 issubdtype를 사용하여 데이터 타입을 확인하고, 적절히 처리
        
        # 만약 입력 텐서가 bytes 타입이라면 as_numpy()로 변환 시 오류 발생 가능
        # 이를 방지하기 위해 issubdtype를 사용하여 조건 처리
        if np.issubdtype(tensor_value.as_numpy().dtype, np.bytes_):
            # bytes 타입 텐서는 바로 디코딩
            return self._np_decoder(tensor_value.as_numpy())
        else:
            # 그 외의 타입은 기존 방식 유지 (첫 번째 요소만 가져옴)
            return self._np_decoder(tensor_value.as_numpy()[0])

    def _np_decoder(self, obj):
        """
        NumPy 객체의 데이터 타입을 확인하고 Python 기본 타입으로 변환합니다.

        Args:
            obj (numpy.ndarray element): 변환할 NumPy 배열의 요소.

        Returns:
            any: 해당 NumPy 요소에 대응하는 Python 기본 타입 (str, int, float, bool).
                 bytes 타입인 경우 UTF-8로 디코딩합니다.
        """
        if isinstance(obj, bytes):
           return obj.decode('utf-8')
        if np.issubdtype(obj, np.integer):
           return int(obj)
        if np.issubdtype(obj, np.floating):
           return round(float(obj), 3)
        if isinstance(obj, np.bool_):
           return bool(obj)        

    def finalize(self):
        """
        모델 실행이 완료된 후 Triton 서버가 종료될 때 한 번 호출되는 함수입니다.
        `finalize` 함수를 구현하는 것은 선택 사항입니다. 이 함수를 통해 모델은
        종료 전에 필요한 모든 정리 작업을 수행할 수 있습니다.
        """
        pass