import os
import torch
import torch.nn.functional as F
from transformers import AutoModelForSequenceClassification, AutoTokenizer

MAX_LENGTH = 128

DEFAULT_MODEL_ID = "FallacyHunter/Fallacy-Hunter-Roberta"


class FallacyClassifier:

    def __init__(self, model_path=None):
        self.model_path = model_path or os.environ.get("ROBERTA_MODEL_ID", DEFAULT_MODEL_ID)
        self.model = None
        self.tokenizer = None
        self._load_model()

    def _load_model(self):
        print(f"Loading RoBERTa tokenizer from {self.model_path}...")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)

        print("Loading RoBERTa model...")
        self.model = AutoModelForSequenceClassification.from_pretrained(
            self.model_path,
            torch_dtype=torch.bfloat16,
            device_map="cpu",
        )
        self.model.eval()
        print("RoBERTa loaded successfully")

    def classify(self, text: str) -> dict:
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            truncation=True,
            padding="max_length",
            max_length=MAX_LENGTH
        )
        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}

        with torch.no_grad():
            logits = self.model(**inputs).logits
            probs = F.softmax(logits.float(), dim=-1)[0]

        id2label = self.model.config.id2label
        probabilities = {id2label[i]: float(probs[i]) for i in range(len(probs))}

        top_idx = int(probs.argmax())
        final_label = id2label[top_idx]
        confidence = float(probs[top_idx])

        return {
            "final_label": final_label,
            "confidence": confidence,
            "probabilities": probabilities
        }