"""Custom handler for HuggingFace Inference Endpoints — TextSight T5 Humanizer"""
from typing import Dict, Any
from transformers import T5ForConditionalGeneration, AutoTokenizer
import torch


class EndpointHandler:
    def __init__(self, path: str = ""):
        # Load tokenizer from HF hub (avoids local spiece.model path issues)
        self.tokenizer = AutoTokenizer.from_pretrained("google-t5/t5-large")
        # Load model weights from the local repo path
        self.model = T5ForConditionalGeneration.from_pretrained(
            path,
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        )
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs", "")
        params = data.get("parameters", {})

        if not inputs:
            return {"error": "No input text provided"}

        # Prefix for T5
        input_text = f"humanize: {inputs}"

        tokens = self.tokenizer(
            input_text,
            return_tensors="pt",
            max_length=512,
            truncation=True,
            padding=True,
        ).to(self.device)

        with torch.no_grad():
            output_ids = self.model.generate(
                **tokens,
                max_new_tokens=params.get("max_new_tokens", 512),
                num_beams=params.get("num_beams", 4),
                temperature=params.get("temperature", 1.1),
                do_sample=True,
                top_p=params.get("top_p", 0.92),
                top_k=params.get("top_k", 50),
                repetition_penalty=params.get("repetition_penalty", 2.5),
                no_repeat_ngram_size=3,
                early_stopping=True,
            )

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