# 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):
        """
        모델이 로드될 때 딱 한 번만 호출됩니다.
        """
        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:
            try:
                # 입력 텐서 직접 가져오기
                self.logger.log_info("입력 텐서 직접 가져오기")
                input_tensor = pb_utils.get_input_tensor_by_name(request, "INPUT")
                self.logger.log_info(f"input_tensor: '{input_tensor}'")

                if input_tensor is None:
                    # 입력이 없는 경우 기본값 처리
                    input_text = "default_input"
                else:
                    # numpy 배열로 변환 후 문자열로 디코딩
                    input_array = input_tensor.as_numpy()
                    if len(input_array) > 0:
                        # bytes인 경우 decode, 아니면 str로 변환
                        first_element = input_array[0]
                        if isinstance(first_element, bytes):
                            input_text = first_element.decode('utf-8')
                        elif isinstance(first_element, np.bytes_):
                            input_text = first_element.decode('utf-8')
                        else:
                            input_text = str(first_element)
                    else:
                        input_text = "empty_input"
                
                self.logger.log_info(f"INPUT 처리됨: {input_text}")

                # 랜덤 문자열 생성
                random_string = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(16))
                output_text = f"Processed: {input_text} -> {random_string}"
                
                self.logger.log_info(f"OUTPUT 생성됨: {output_text}")

                # 출력 텐서 생성 - bytes 형태로 직접 생성
                output_bytes = output_text.encode('utf-8')
                output_tensor = pb_utils.Tensor("OUTPUT", np.array([output_bytes], dtype=np.object_))

                # 응답 생성
                response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
                responses.append(response)

            except Exception as e:
                # 오류 발생시 오류 응답 생성
                self.logger.log_error(f"실행 중 오류 발생: {str(e)}")
                error_msg = f"Error: {str(e)}"
                output_bytes = error_msg.encode('utf-8')
                output_tensor = pb_utils.Tensor("OUTPUT", np.array([output_bytes], dtype=np.object_))
                response = pb_utils.InferenceResponse(output_tensors=[output_tensor])
                responses.append(response)

        return responses

    def finalize(self):
        """
        모델 종료시 정리 작업
        """
        self.logger.log_info(f"'{self.model_name}' 모델 종료")
        pass