from transformers import (
    AutoConfig,
    AutoTokenizer,
    AutoModelForSequenceClassification,
)
import torch
from typing import Dict, Any, List


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

        config = AutoConfig.from_pretrained(path)
        config.num_labels = 1

        self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)

        self.model = AutoModelForSequenceClassification.from_pretrained(
            path,
            config=config,
            ignore_mismatched_sizes=True,
        )
        self.model.to(self.device)
        self.model.eval()

        self.threshold = float(getattr(self.model.config, "threshold_F1", 0.5))
        self.num_labels = int(getattr(self.model.config, "num_labels", 1))

    def _predict_probs(self, inputs_text: List[str]) -> List[float]:
        encoded = self.tokenizer(
            inputs_text,
            return_tensors="pt",
            truncation=True,
            padding=True,
            max_length=512,
        )
        encoded = {k: v.to(self.device) for k, v in encoded.items()}

        with torch.no_grad():
            logits = self.model(**encoded).logits

        if self.num_labels == 1 or logits.shape[-1] == 1:
            logits = logits.squeeze(-1)
            probs = torch.sigmoid(logits).detach().cpu().tolist()
        else:
            probs = torch.softmax(logits, dim=-1)[:, 1].detach().cpu().tolist()

        if isinstance(probs, float):
            probs = [probs]

        return probs

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

        if isinstance(inputs_text, str):
            inputs_text = [inputs_text]
        elif not isinstance(inputs_text, list):
            raise ValueError("`inputs` must be a string or a list of strings.")

        inputs_text = [
            str(x).strip()
            for x in inputs_text
            if x is not None and str(x).strip() != ""
        ]

        if not inputs_text:
            return []

        probs = self._predict_probs(inputs_text)

        return [
            {
                "text": text,
                "label": "CAUSAL" if prob >= self.threshold else "NON-CAUSAL",
                "score": round(float(prob), 6),
                "p_causal": round(float(prob), 6),
                "threshold": round(float(self.threshold), 6),
            }
            for text, prob in zip(inputs_text, probs)
        ]