# modify 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):
        responses = []

        for request in requests:
            # INPUT
            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)
            )
            self.logger.log_info(f"OUTPUT 출력:\n{random_string}")

            # Triton이 요구하는 shape [1] STRING 출력
            output_tensor = pb_utils.Tensor(
                "OUTPUT",
                np.array([random_string.encode("utf-8")], dtype=np.object_)
            )

            responses.append(pb_utils.InferenceResponse(
                output_tensors=[output_tensor]
            ))

        return responses

    def _get_input_value(self, request, input_name: str, default=None):
        """
        Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다.

        Args:
            request (pb_utils.InferenceRequest): Triton 추론 요청 객체.
            input_name (str): 가져올 입력 텐서의 이름.
            default (any, optional): 입력 텐서가 없을 경우 반환할 기본값. Defaults to None.

        Returns:
            any: 디코딩된 입력 텐서의 값. 텐서가 없으면 기본값을 반환합니다.
        """
        tensor_value = pb_utils.get_input_tensor_by_name(request, input_name)

        if tensor_value is None:
            return default

        return self._np_decoder(tensor_value.as_numpy()[0])

    def _np_decoder(self, obj):
        if isinstance(obj, (bytes, np.bytes_)):
            return obj.decode('utf-8')
        if isinstance(obj, str):
            return obj
        if np.issubdtype(type(obj), np.integer):
            return int(obj)
        if np.issubdtype(type(obj), np.floating):
            return round(float(obj), 3)
        if isinstance(obj, (np.bool_, bool)):
            return bool(obj)
        return obj       

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