from typing import Any, Dict

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


class EndpointHandler:
    def __init__(self, path=""):
        base_model = AutoModelForCausalLM.from_pretrained(
            "microsoft/Phi-3-mini-4k-instruct",
            device_map="auto",
            torch_dtype="auto",
        )
        self.model = PeftModel.from_pretrained(base_model, path)
        self.tokenizer = AutoTokenizer.from_pretrained(
            "microsoft/Phi-3-mini-4k-instruct"
        )
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
        question = data.get("inputs")
        if not isinstance(question, str) or not question.strip():
            raise ValueError("Envía una pregunta de texto en el campo 'inputs'.")

        parameters = data.get("parameters") or {}
        max_new_tokens = int(parameters.get("max_new_tokens", 256))
        max_new_tokens = max(1, min(max_new_tokens, 512))

        messages = [{"role": "user", "content": question}]
        input_ids = self.tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        )
        input_ids = input_ids.to(self.model.get_input_embeddings().weight.device)

        with torch.inference_mode():
            generated = self.model.generate(
                input_ids=input_ids,
                max_new_tokens=max_new_tokens,
                do_sample=False,
                eos_token_id=self.tokenizer.eos_token_id,
            )

        new_tokens = generated[0, input_ids.shape[-1]:]
        answer = self.tokenizer.decode(
            new_tokens,
            skip_special_tokens=True,
        ).strip()

        return {"generated_text": answer}