from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

MODEL_NAME = "teapotai/tinyteapot"
MAX_INPUT_TOKENS = 512


class EndpointHandler:
    def __init__(self, path: str = ""):
        # EXACTLY like your snippet: load by repo id from Hugging Face Hub
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
        self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

        self.model.eval()
        self.device = torch.device("cpu")
        self.model.to(self.device)

        self.system_prompt = (
            "You are Teapot, an open-source AI assistant optimized for low-end devices, "
            "providing short, accurate responses without hallucinating while excelling at "
            "information extraction and text summarization. "
            "If the context does not answer the question, reply exactly: "
            "'I am sorry but I don't have any information on that'."
        )

    @torch.inference_mode()
    def __call__(self, data):
        inputs = data.get("inputs")
        if inputs is None:
            raise ValueError("Missing required field 'inputs'.")

        # Support:
        # 1) inputs is a full prompt string (closest to your snippet)
        # 2) inputs is {"context": ..., "question": ...}
        if isinstance(inputs, str):
            prompt = inputs
        elif isinstance(inputs, dict):
            context = inputs.get("context", "")
            question = inputs.get("question", "")
            prompt = f"{context}\n{self.system_prompt}\n{question}\n"
        else:
            raise ValueError("inputs must be a string OR an object with 'context' and 'question'.")

        # Match your code path
        enc = self.tokenizer(prompt, return_tensors="pt")
        input_ids = enc["input_ids"]
        attention_mask = enc["attention_mask"]

        # Keep only the most recent tokens if too long
        if input_ids.shape[1] > MAX_INPUT_TOKENS:
            input_ids = input_ids[:, -MAX_INPUT_TOKENS:]
            attention_mask = attention_mask[:, -MAX_INPUT_TOKENS:]

        input_ids = input_ids.to(self.device)
        attention_mask = attention_mask.to(self.device)

        outputs = self.model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            do_sample=False,
        )

        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
        return {"generated_text": answer}