"""
Custom handler for BERT-OJA-SkillLess on HF Inference Endpoints.
Uses ONNX Runtime with CUDA for 2-4x faster inference.
"""
from typing import Dict, List, Any
import numpy as np
from transformers import AutoTokenizer


class EndpointHandler:
    def __init__(self, path=""):
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.batch_size = 1024

        try:
            from optimum.onnxruntime import ORTModelForSequenceClassification
            self.model = ORTModelForSequenceClassification.from_pretrained(
                path, export=True, provider="CUDAExecutionProvider",
            )
            self._use_ort = True
            print(f"[handler] Loaded ONNX model on CUDA (batch_size={self.batch_size})")
        except Exception as e:
            print(f"[handler] ONNX failed ({e}), falling back to PyTorch FP16")
            import torch
            from transformers import AutoModelForSequenceClassification
            self.model = AutoModelForSequenceClassification.from_pretrained(path)
            self.model.eval()
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            if self.device == "cuda":
                self.model = self.model.to(self.device).half()
            self._use_ort = False
            self._torch = torch

    def __call__(self, data: Dict[str, Any]) -> List[List[Dict[str, float]]]:
        inputs = data.get("inputs", data.get("input", ""))
        if isinstance(inputs, str):
            inputs = [inputs]

        all_results = []
        for i in range(0, len(inputs), self.batch_size):
            batch = inputs[i : i + self.batch_size]
            encoded = self.tokenizer(
                batch, padding=True, truncation=True,
                max_length=128, return_tensors="pt" if not self._use_ort else "np",
            )

            if self._use_ort:
                logits = self.model(**{k: v for k, v in encoded.items()}).logits
                if hasattr(logits, 'numpy'):
                    logits = logits.numpy()
                exp = np.exp(logits - logits.max(axis=-1, keepdims=True))
                probs = exp / exp.sum(axis=-1, keepdims=True)
                for j in range(len(batch)):
                    all_results.append([
                        {"label": "LABEL_0", "score": round(float(probs[j][0]), 6)},
                        {"label": "LABEL_1", "score": round(float(probs[j][1]), 6)},
                    ])
            else:
                torch = self._torch
                encoded = {k: v.to(self.device) for k, v in encoded.items()}
                with torch.no_grad():
                    logits = self.model(**encoded).logits
                    probs = torch.softmax(logits, dim=-1)
                for j in range(len(batch)):
                    all_results.append([
                        {"label": "LABEL_0", "score": round(probs[j][0].item(), 6)},
                        {"label": "LABEL_1", "score": round(probs[j][1].item(), 6)},
                    ])

        return all_results