import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

def EndpointHandler(path=""):
    """
    Inference Endpoints에서 사용할 핸들러 함수
    """
    class Handler:
        def __init__(self, path=""):
            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            self.base_model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
            self.adapter_path = path or "./tinyllama-qa-news"

            # ✅ 모델 로딩 (GPU)
            base_model = AutoModelForCausalLM.from_pretrained(
                self.base_model_id,
                torch_dtype=torch.float16
            )
            self.model = PeftModel.from_pretrained(base_model, self.adapter_path)
            self.model.to(self.device)
            self.model.eval()

            self.tokenizer = AutoTokenizer.from_pretrained(self.base_model_id)
            self.tokenizer.pad_token = self.tokenizer.eos_token
            self.tokenizer.padding_side = "right"

        def generate_response(self, user_input: str):
            messages = [
                {
                    "role": "system",
                    "content": "당신은 공감적이고 이해심 깊은 감정 상담 챗봇입니다. 사용자의 감정과 생각을 이해하고 공감하며, 적절한 조언을 제공합니다."
                },
                {"role": "user", "content": user_input}
            ]

            prompt = self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True,
                return_tensors=None
            )

            inputs = self.tokenizer(prompt, return_tensors="pt", padding=True)
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=256,
                    temperature=0.8,
                    top_p=0.95,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id,
                    repetition_penalty=1.1,
                    length_penalty=1.0,
                    num_return_sequences=1,
                    num_beams=1,
                    early_stopping=False
                )

            decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)

            if "<|assistant|>" in decoded:
                response = decoded.split("<|assistant|>")[-1].strip()
            else:
                response = decoded.split(prompt)[-1].strip()

            if not response:
                response = "감정을 이해해 주셔서 감사합니다. 더 이야기하고 싶으신 게 있으신가요?"

            return response

        def __call__(self, data):
            if not data or "inputs" not in data:
                return {"error": "입력 데이터가 없습니다."}

            user_input = data["inputs"]
            response = self.generate_response(user_input)
            return {"response": response}
    
    return Handler(path)