import torch
import json
import numpy as np
import triton_python_backend_utils as pb_utils

from transformers import AutoModel, AutoTokenizer, GenerationConfig
from peft import PeftModel, PeftConfig


class TritonPythonModel:
    def initialize(self, args):
        """
        모델이 로드될 때 딱 한 번만 호출됩니다.
        """
        self.logger = pb_utils.Logger

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

        self.model_name = args["model_name"]
        self.base_model_path = self._get_config_parameter("base_model_path")
        
        self.logger.log_info(f"base_model_path: {self.base_model_path}")
        

        self.load_model()

    def load_model(self):
        torch_dtype = torch.float16

        # AutoModel로 모델 로드
        self.model = AutoModel.from_pretrained(
            pretrained_model_name_or_path=self.base_model_path,
            local_files_only=True,
            trust_remote_code=True
        )
        
        # 모델을 평가 모드로 설정
        self.model.eval()

        self.tokenizer = AutoTokenizer.from_pretrained(
                                        self.base_model_path,
                                        trust_remote_code=True,
                                        add_eos_token=True
                                        )

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

    def execute(self, requests):
        """
        Triton이 각 추론 요청에 대해 호출하는 실행 함수입니다.
        Generation 대신 Embedding 생성을 수행하도록 수정합니다.
        """
        responses = []

        # 각 추론 요청을 순회하며 처리합니다.
        for request in requests:
            # Triton 입력 파싱: 텍스트 입력만 처리합니다.
            input_text = self._get_input_value(request, "text_input")

            # CodeSage는 대화 형식이 아닌 일반 텍스트 (코드)를 입력으로 받으므로 
            # JSON 파싱 로직과 Chat 템플릿 로직을 제거합니다.
            text = input_text
            self.logger.log_info(f"입력 text 출력:\n{text}")
            
            # 입력 텍스트를 토큰화합니다.
            # add_eos_token=True가 load_model에서 설정되었으므로 토큰화 시 자동으로 추가됩니다.
            inputs = self.tokenizer(
                        text,
                        return_tensors="pt").to(device=self.model.device)

            input_ids = inputs["input_ids"]
            attention_mask = inputs["attention_mask"]
            
            # # CodeSage는 텍스트 생성이 아닌 임베딩 생성을 수행합니다.
            # gened = self.model.generate(...)

            # **임베딩 생성**
            # CodeSage 모델은 (임베딩, 히든 스테이트, 어텐션)을 반환하며, 첫 번째 요소가 임베딩입니다.
            with torch.no_grad():
                # inputs에는 input_ids와 attention_mask가 모두 포함되어 전달됩니다.
                outputs = self.model(**inputs) 
            
            # outputs[0]에는 임베딩 벡터가 포함되어 있습니다.
            # 임베딩은 일반적으로 첫 번째 토큰 (CLS 토큰 또는 문맥 임베딩)을 사용합니다.
            # CodeSage의 경우, 모델 카드 예시를 보면 outputs[0] 전체를 사용합니다.
            # 여기서는 [batch_size, sequence_length, hidden_size] 형태의 임베딩 중 첫 번째 토큰 임베딩을 사용합니다.
            # 임베딩 사용법은 모델의 목적에 따라 달라질 수 있습니다. CodeSage는 주로 전체 시퀀스 임베딩을 사용합니다.
            # 여기서는 예시와 같이 첫 번째 요소 (last_hidden_state)를 가져옵니다.
            # 임베딩 크기: [1, seq_len, hidden_size]
            embeddings = outputs[0] 
            
            # 임베딩을 NumPy 배열로 변환
            # CPU로 옮기고, NumPy로 변환
            embedding_np = embeddings.squeeze().cpu().numpy()
            
            # 출력 텐서 생성 (데이터 타입은 float32 또는 float16이 적합)
            # CodeSage는 단일 문장 입력만 처리하므로, 배치 차원 없이 [seq_len, hidden_size]로 가정합니다.
            # 실제 사용 목적에 따라 풀링 로직을 추가하여 [hidden_size] 벡터로 만들 수도 있습니다.
            output_tensor = pb_utils.Tensor("embedding_output", embedding_np.astype(np.float32))

            self.logger.log_info(f"모델이 생성한 임베딩 Shape:\n{embedding_np.shape}")

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

        return responses
    

    def _get_config_parameter(self, parameter_name):
        """
        모델 설정(config.pbtxt)에서 특정 파라미터의 문자열 값을 가져옵니다.
        """
        self.parameters = self.model_config.get('parameters', {})
        parameter_dict = self.parameters.get(parameter_name)

        if isinstance(parameter_dict, dict) and 'string_value' in parameter_dict:
            return parameter_dict['string_value']

        return None

    def _get_input_value(self, request, input_name: str, default=None):
        """
        Triton 추론 요청에서 특정 이름의 입력 텐서 값을 가져옵니다.
        """
        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):
        """
        NumPy 객체의 데이터 타입을 확인하고 Python 기본 타입으로 변환합니다.
        """
        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 서버가 종료될 때 한 번 호출되는 함수입니다.
        """
        pass