import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import threading
from typing import Dict, Any

class DebertaClassifier:
    _instance = None
    _lock = threading.Lock()

    def __new__(cls):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super(DebertaClassifier, cls).__new__(cls)
                cls._instance._initialized = False
        return cls._instance

    def initialize(self):
        if self._initialized:
            return
        
        # Load model and tokenizer
        self.model_name = "deepset/deberta-v3-base-injection"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
        self.device = "cpu"  # Force CPU since it's a lightweight laptop configuration
        self.model.to(self.device)
        self._initialized = True

    def predict(self, text: str) -> Dict[str, Any]:
        """
        Classifies whether input text is a prompt injection.
        Returns safety report with probability score.
        """
        if not self._initialized:
            self.initialize()

        # Tokenize and evaluate
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        with torch.no_grad():
            outputs = self.model(**inputs)
            logits = outputs.logits
            probabilities = torch.softmax(logits, dim=1)[0]
        
        # Class labels: 0: SAFE, 1: INJECTION
        safe_prob = float(probabilities[0])
        injection_prob = float(probabilities[1])

        # Return dict report
        is_safe = safe_prob > 0.5
        return {
            "safe": is_safe,
            "risk_score": round(injection_prob, 4),
            "category": "safe" if is_safe else "prompt_injection"
        }

# Global Singleton access instance
classifier = DebertaClassifier()