"""
PhishShield — RAG Engine (ML Classifier)
Uses fine-tuned BERT model for email phishing detection
"""

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch


class RAGEngine:
    """
    Email phishing classifier using JellyPhish BERT model.
    Lightweight and designed for phishing detection.
    """

    def __init__(self, model_name: str = "RamzyBakir/jellyphish-bert-base-mail"):
        print(f"[RAG] Loading {model_name}...")
        self.model_name = model_name
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.eval()
        print("[RAG] ✅ Model loaded successfully")

    def query(self, text: str) -> tuple[str, float]:
        """
        Return (classification_result, confidence_score)
        - Returns "Phishing" or "Legitimate"
        - Confidence score 0.0 - 1.0
        """
        if not text.strip():
            return "No content", 0.0

        # Truncate to max length
        if len(text) > 5000:
            text = text[:5000]

        inputs = self.tokenizer(
            text,
            truncation=True,
            padding=True,
            max_length=512,
            return_tensors="pt"
        )
        
        with torch.no_grad():
            outputs = self.model(**inputs)
            probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
            # Class 0 = Legitimate, Class 1 = Phishing
            confidence = probabilities[0][1].item()
            
        predicted_class = torch.argmax(outputs.logits, dim=-1).item()
        label = "Phishing" if predicted_class == 1 else "Legitimate"
        
        return label, confidence

    def doc_count(self) -> int:
        return 1

    def add_document(self, text: str):
        """Placeholder for compatibility"""
        pass