import os
import torch
import numpy as np
import joblib

# Hugging Face가 제공하는 기본 핸들러 클래스를 가져옵니다.
from text_generation_server.models.custom_handler import BaseHandler

class TidePredictionHandler(BaseHandler):
    """
    TimeXer 모델을 위한 커스텀 핸들러
    """
    def __init__(self, model, tokenizer):
        # 이 함수는 서버가 시작될 때 딱 한 번 실행됩니다.
        # 모델 초기화, 스케일러 로딩 등 준비 작업을 여기서 합니다.
        super().__init__(model, tokenizer)
        
        # 1. 스케일러를 불러와서 self.scaler에 저장합니다.
        #    경로는 저장소 내의 실제 파일 위치와 같아야 합니다.
        scaler_path = os.path.join(os.getcwd(), 'checkpoints', 'scaler.gz')
        self.scaler = joblib.load(scaler_path)
        
        # 2. 모델을 평가 모드로 설정합니다.
        self.model.eval()
        
        # 3. 모델의 설정값들을 self.model.args 처럼 접근할 수 있도록 저장해두면 편리합니다.
        #    이 부분은 TimeXer 모델의 구조에 따라 필요 없을 수도 있습니다.
        #    예: self.seq_len = self.model.seq_len

    def __call__(self, inputs, **kwargs):
        # 이 함수는 API 예측 요청이 올 때마다 실행됩니다.
        # 실제 예측 로직이 들어가는 부분입니다.

        # 1. 입력 데이터 파싱
        # inputs는 보통 리스트 형태의 텍스트 또는 바이트 데이터로 들어옵니다.
        # JSON 형식으로 입력을 받으려면 추가적인 처리가 필요할 수 있습니다.
        # 여기서는 간단히 inputs가 숫자 리스트 문자열이라고 가정합니다.
        # 예: "500.1, 502.3, ..., 498.7"
        
        # 문자열을 숫자 리스트로 변환
        try:
            # 입력 데이터를 파싱하는 가장 좋은 방법은 JSON을 사용하는 것입니다.
            # 예: `json.loads(inputs[0])`
            # 여기서는 간단한 예시를 위해 split을 사용합니다.
            input_list = [float(i) for i in inputs[0].split(',')]
            seq_len = 144 # 이 값은 실제 모델의 입력 길이와 일치해야 합니다.

            if len(input_list) != seq_len:
                 raise ValueError(f"Input must have {seq_len} items.")

        except Exception as e:
            # 오류 발생 시 에러 메시지를 반환합니다.
            return {"error": f"Invalid input format: {str(e)}"}, 400

        # 2. 데이터를 모델 입력 형식(Tensor)으로 변환
        input_array = np.array(input_list).reshape(-1, 1)
        scaled_input = self.scaler.transform(input_array)
        input_tensor = torch.from_numpy(scaled_input).float().unsqueeze(0).to(self.model.device)

        # 3. 모델 예측 실행
        with torch.no_grad():
            # TimeXer 모델의 forward 함수에 필요한 모든 인자를 전달해야 합니다.
            # 예: outputs = self.model(batch_x=input_tensor, batch_x_mark=...)
            # 이 부분은 모델의 실제 코드를 보고 채워야 합니다.
            # 여기서는 input_tensor만 필요하다고 가정합니다.
            outputs = self.model(input_tensor)

        # 4. 예측 결과를 후처리하고 원래 스케일로 복원
        prediction_scaled = outputs.detach().cpu().numpy().squeeze()
        prediction = self.scaler.inverse_transform(prediction_scaled.reshape(-1, 1))
        
        # 5. 최종 결과를 리스트 형태로 반환
        # Hugging Face 핸들러는 보통 텍스트나 바이트를 반환해야 합니다.
        # 결과를 JSON 문자열로 만들어 반환하는 것이 일반적입니다.
        import json
        result_str = json.dumps({"prediction": prediction.flatten().tolist()})

        return [result_str]