app.py

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download
import gradio as gr
import requests
import re
from urllib.parse import urlparse
from bs4 import BeautifulSoup
import time
import joblib



# --- import your architecture ---
# Make sure this file is in the repo (e.g., models/deberta_lstm_classifier.py)
# and update the import path accordingly.
from model import DeBERTaLSTMClassifier  # <-- your class

# --------- Config ----------
REPO_ID = "khoa-done/phishing-detector"       # HF repo that holds the checkpoint
CKPT_NAME = "deberta_lstm_checkpoint.pt"      # the .pt file name
MODEL_NAME = "microsoft/deberta-base"         # base tokenizer/backbone
LABELS = ["benign", "phishing"]               # adjust to your classes

# If your checkpoint contains hyperparams, you can fetch them like:
# checkpoint.get("config") or checkpoint.get("model_args")
# and pass into DeBERTaLSTMClassifier(**model_args)

# --------- Load model/tokenizer once (global) ----------
device = "cuda" if torch.cuda.is_available() else "cpu"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

ckpt_path = hf_hub_download(repo_id=REPO_ID, filename=CKPT_NAME)
checkpoint = torch.load(ckpt_path, map_location=device)

# If you saved hyperparams in the checkpoint, use them:
model_args = checkpoint.get("model_args", {})  # e.g., {"lstm_hidden":256, "num_labels":2, ...}
model = DeBERTaLSTMClassifier(**model_args)

# Load state dict and handle missing attention layer for older models
try:
    model.load_state_dict(checkpoint["model_state_dict"])
except RuntimeError as e:
    if "attention" in str(e):
        # Old model without attention layer - initialize attention layer and load partial state
        state_dict = checkpoint["model_state_dict"]
        model_dict = model.state_dict()
        # Filter out attention layer parameters
        filtered_dict = {k: v for k, v in state_dict.items() if "attention" not in k}
        model_dict.update(filtered_dict)
        model.load_state_dict(model_dict)
        print("Loaded model without attention layer, using newly initialized attention weights")
    else:
        raise e

model.to(device).eval()

# --------- Load BERT model/tokenizer from Hugging Face Hub ----------
BERT_MODEL_PATH = "th1enq/bert_checkpoint"  # Use Hugging Face Hub model
bert_tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_PATH)
bert_model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL_PATH)
bert_model.to(device).eval()

# --------- Helper functions ----------
def is_url(text):
    """Check if text is a URL"""
    url_pattern = re.compile(
        r'^https?://'  # http:// or https://
        r'(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+[A-Z]{2,6}\.?|'  # domain...
        r'localhost|'  # localhost...
        r'\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})'  # ...or ip
        r'(?::\d+)?'  # optional port
        r'(?:/?|[/?]\S+)$', re.IGNORECASE)
    return url_pattern.match(text) is not None

def fetch_html_content(url, timeout=10):
    """Fetch HTML content from URL"""
    try:
        headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
        }
        response = requests.get(url, headers=headers, timeout=timeout, verify=False)
        response.raise_for_status()
        return response.text, response.status_code
    except requests.exceptions.RequestException as e:
        return None, f"Request error: {str(e)}"
    except Exception as e:
        return None, f"General error: {str(e)}"

def predict_single_text(text, text_type="text"):
    """Predict for a single text input"""
    # Tokenize
    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=256
    )
    # DeBERTa typically doesn't use token_type_ids
    inputs.pop("token_type_ids", None)
    # Move to device
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        try:
            # Try to get predictions with attention weights
            result = model(**inputs, return_attention=True)
            if isinstance(result, tuple) and len(result) == 3:
                logits, attention_weights, deberta_attentions = result
                has_attention = True
            else:
                logits = result
                has_attention = False
        except TypeError:
            # Fallback for older model without return_attention parameter
            logits = model(**inputs)
            has_attention = False
            
        probs = F.softmax(logits, dim=-1).squeeze(0).tolist()

    # Get tokens for visualization
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'].squeeze(0).tolist())
    
    return probs, tokens, has_attention, attention_weights if has_attention else None

def combine_predictions(url_probs, html_probs, url_weight=0.3, html_weight=0.7):
    """Combine URL and HTML content predictions"""
    combined_probs = [
        url_weight * url_probs[0] + html_weight * html_probs[0],  # benign
        url_weight * url_probs[1] + html_weight * html_probs[1]   # phishing
    ]
    return combined_probs

# --------- Inference function ----------
def predict_fn(text: str):
    if not text or not text.strip():
        return {"error": "Please enter a URL or text."}, ""

    # Check if input is URL
    if is_url(text.strip()):
        # Process URL
        url = text.strip()
        
        # Get prediction for URL itself
        url_probs, url_tokens, url_has_attention, url_attention = predict_single_text(url, "URL")
        
        # Try to fetch HTML content
        html_content, status = fetch_html_content(url)
        
        if html_content:
            # Get prediction for HTML content
            html_probs, html_tokens, html_has_attention, html_attention = predict_single_text(html_content, "HTML")
            
            # Combine predictions
            combined_probs = combine_predictions(url_probs, html_probs)
            
            # Use combined probabilities but show analysis for both
            probs = combined_probs
            tokens = url_tokens + ["[SEP]"] + html_tokens[:50]  # Limit HTML tokens for display
            has_attention = url_has_attention or html_has_attention
            attention_weights = url_attention if url_has_attention else html_attention
            
            analysis_type = "Combined URL + HTML Analysis"
            fetch_status = f"✅ Successfully fetched HTML content (Status: {status})"
            
        else:
            # Fallback to URL-only analysis
            probs = url_probs
            tokens = url_tokens
            has_attention = url_has_attention
            attention_weights = url_attention
            
            analysis_type = "URL-only Analysis"
            fetch_status = f"⚠️ Could not fetch HTML content: {status}"
    else:
        # Process as regular text
        probs, tokens, has_attention, attention_weights = predict_single_text(text, "text")
        analysis_type = "Text Analysis"
        fetch_status = ""

    # Get tokens for visualization
    
    # Create detailed analysis
    predicted_class = "phishing" if probs[1] > probs[0] else "benign"
    confidence = max(probs)
    
    detailed_analysis = f"""
<div style="font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; background: #1e1e1e; padding: 20px; border-radius: 15px;">
<div style="background: linear-gradient(135deg, {'#8b0000' if predicted_class == 'phishing' else '#006400'} 0%, {'#dc143c' if predicted_class == 'phishing' else '#228b22'} 100%); padding: 25px; border-radius: 20px; color: white; text-align: center; margin-bottom: 20px; box-shadow: 0 8px 32px rgba(0,0,0,0.5); border: 2px solid {'#ff4444' if predicted_class == 'phishing' else '#44ff44'};">
    <h2 style="margin: 0 0 10px 0; font-size: 28px; color: white;">🔍 {analysis_type}</h2>
    <div style="font-size: 36px; font-weight: bold; margin: 10px 0; color: white;">
        {predicted_class.upper()}
    </div>
    <div style="font-size: 18px; color: #f0f0f0;">
        Confidence: {confidence:.1%}
    </div>
    <div style="margin-top: 15px; font-size: 14px; color: #e0e0e0;">
        {'This appears to be a phishing attempt!' if predicted_class == 'phishing' else '✅ This appears to be legitimate content.'}
    </div>
</div>
"""

    if fetch_status:
        detailed_analysis += f"""
<div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
    <strong>Fetch Status:</strong> {fetch_status}
</div>
"""
    
    if has_attention and attention_weights is not None:
        attention_scores = attention_weights.squeeze(0).tolist()
        
        token_analysis = []
        for i, (token, score) in enumerate(zip(tokens, attention_scores)):
            # More lenient filtering - include more tokens for text analysis
            if token not in ['[CLS]', '[SEP]', '[PAD]', '<s>', '</s>'] and len(token.strip()) > 0 and score > 0.005:  
                clean_token = token.replace('▁', '').replace('Ġ', '').strip()  # Handle different tokenizer prefixes
                if clean_token:  # Only add if token has content after cleaning
                    token_analysis.append({
                        'token': clean_token,
                        'importance': score,
                        'position': i
                    })
        
        # Sort by importance
        token_analysis.sort(key=lambda x: x['importance'], reverse=True)
        
        detailed_analysis += f"""
## Top important tokens:
<div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
    <strong>Analysis Info:</strong> Found {len(token_analysis)} important tokens out of {len(tokens)} total tokens
</div>
<div style="font-family: Arial, sans-serif;">
"""
        
        for i, token_info in enumerate(token_analysis[:10]):  # Top 10 tokens
            bar_width = int(token_info['importance'] * 100)
            color = "#ff4444" if predicted_class == "phishing" else "#44ff44"
            
            detailed_analysis += f"""
<div style="margin: 8px 0; display: flex; align-items: center; background: #2d2d2d; padding: 8px; border-radius: 8px; border-left: 4px solid {color};">
    <div style="width: 30px; text-align: right; margin-right: 10px; font-weight: bold; color: #ffffff;">
        {i+1}.
    </div>
    <div style="width: 120px; margin-right: 10px; font-weight: bold; color: #e0e0e0; text-align: right;">
        {token_info['token']}
    </div>
    <div style="width: 300px; background-color: #404040; border-radius: 10px; overflow: hidden; margin-right: 10px; border: 1px solid #555;">
        <div style="width: {bar_width}%; background-color: {color}; height: 20px; border-radius: 10px; transition: width 0.3s ease;"></div>
    </div>
    <div style="color: #cccccc; font-size: 12px; font-weight: bold;">
        {token_info['importance']:.1%}
    </div>
</div>
"""
        
        detailed_analysis += "</div>\n"
        
        detailed_analysis += f"""
## Detailed analysis:
<div style="font-family: Arial, sans-serif; background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
    <h3 style="margin: 0 0 15px 0; color: white;">Statistical Overview</h3>
    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 15px;">
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Total tokens</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold, color: white;">{len([t for t in token_analysis if t['importance'] > 0.05])}</div>
            <div style="font-size: 14px, color: #e0e0e0;">High impact tokens (>5%)</div>
        </div>
    </div>
</div>
<div style="font-family: Arial, sans-serif; margin: 15px 0; background: #2d2d2d; padding: 20px; border-radius: 15px; border: 1px solid #555;">
    <h3 style="color: #ffffff; margin-bottom: 15px;"> Prediction Confidence</h3>
    <div style="display: flex; justify-content: space-between; margin-bottom: 10px;">
        <span style="font-weight: bold; color: #ff4444;">Phishing</span>
        <span style="font-weight: bold; color: #44ff44;">Benign</span>
    </div>
    <div style="width: 100%; background-color: #404040; border-radius: 25px; overflow: hidden; height: 30px; border: 1px solid #666;">
        <div style="width: {probs[1]*100:.1f}%; background: linear-gradient(90deg, #ff4444 0%, #ff6666 100%); height: 100%; display: flex; align-items: center; justify-content: center; color: white; font-weight: bold; font-size: 14px;">
            {probs[1]:.1%}
        </div>
    </div>
    <div style="margin-top: 10px; text-align: center; color: #cccccc; font-size: 14px;">
        Benign: {probs[0]:.1%}
    </div>
</div>
"""
    else:
        # Fallback analysis without attention weights
        detailed_analysis += f"""
<div style="background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
    <h3 style="margin: 0 0 15px 0; color: white;">Basic Analysis</h3>
    <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 15px;">
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[1]:.1%}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Phishing</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[0]:.1%}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Benign</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Tokens</div>
        </div>
    </div>
</div>
<div style="background: #2d2d2d; padding: 20px; border-radius: 15px; margin: 15px 0; border: 1px solid #555;">
    <h3 style="color: #ffffff; margin: 0 0 15px 0;">🔤 Tokens in text:</h3>
    <div style="display: flex; flex-wrap: wrap; gap: 8px;">""" + ''.join([f'<span style="background: #404040; color: #64b5f6; padding: 4px 8px; border-radius: 15px; font-size: 12px; border: 1px solid #666;">{token.replace("▁", "")}</span>' for token in tokens if token not in ['[CLS]', '[SEP]', '[PAD]']]) + f"""</div>
    <div style="margin-top: 15px; padding: 10px; background: #3d2914; border-radius: 8px; border-left: 4px solid #ff9800;">
        <strong style="color: #ffcc02;">Debug info:</strong> <span style="color: #e0e0e0;">Found {len(tokens)} total tokens, {len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])} content tokens</span>
    </div>
</div>
<div style="background: #3d2914; padding: 15px; border-radius: 10px; border-left: 4px solid #ff9800; margin: 15px 0;">
    <p style="margin: 0; color: #ffcc02; font-size: 14px;">
        <strong>Note:</strong> Detailed attention weights analysis is not available for the current model.
    </p>
</div>
"""

    # Build label->prob mapping for Gradio Label output
    if len(LABELS) == len(probs):
        prediction_result = {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
    else:
        prediction_result = {f"class_{i}": float(p) for i, p in enumerate(probs)}
    
    return prediction_result, detailed_analysis

# --------- BERT Model Functions ----------
def predict_bert_single_text(text, text_type="text"):
    """Predict for a single text input using BERT."""
    # Tokenize
    inputs = bert_tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        padding=True,
        max_length=512
    )
    # Move to device
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = bert_model(**inputs, output_attentions=True)
        logits = outputs.logits
        
    probs = F.softmax(logits, dim=-1).squeeze(0).tolist()
    
    # Get tokens for visualization
    tokens = bert_tokenizer.convert_ids_to_tokens(inputs['input_ids'].squeeze(0).tolist())
    
    # Get attention weights (use last layer, first head as approximation)
    attention_weights = None
    has_attention = False
    if hasattr(outputs, 'attentions') and outputs.attentions is not None:
        # Average attention across all heads in the last layer
        last_layer_attention = outputs.attentions[-1]  # Last layer
        attention_weights = last_layer_attention.mean(dim=1).squeeze(0)  # Average across heads
        # Use attention to [CLS] token as importance scores
        attention_weights = attention_weights[0]  # [CLS] token attention to all tokens
        has_attention = True
    
    return probs, tokens, has_attention, attention_weights

def predict_bert_interface_fn(text: str):
    """Gradio interface function for BERT model."""
    if not text or not text.strip():
        return {"error": "Please enter a URL or text."}, ""

    # Check if input is URL
    if is_url(text.strip()):
        # Process URL
        url = text.strip()
        
        # Get prediction for URL itself
        url_probs, url_tokens, url_has_attention, url_attention = predict_bert_single_text(url, "URL")
        
        # Try to fetch HTML content
        html_content, status = fetch_html_content(url)
        
        if html_content:
            # Get prediction for HTML content
            html_probs, html_tokens, html_has_attention, html_attention = predict_bert_single_text(html_content, "HTML")
            
            # Combine predictions
            combined_probs = combine_predictions(url_probs, html_probs)
            
            # Use combined probabilities but show analysis for both
            probs = combined_probs
            tokens = url_tokens + ["[SEP]"] + html_tokens[:50]  # Limit HTML tokens for display
            has_attention = url_has_attention or html_has_attention
            attention_weights = url_attention if url_has_attention else html_attention
            
            analysis_type = "Combined URL + HTML BERT Analysis"
            fetch_status = f"✅ Successfully fetched HTML content (Status: {status})"
            
        else:
            # Fallback to URL-only analysis
            probs = url_probs
            tokens = url_tokens
            has_attention = url_has_attention
            attention_weights = url_attention
            
            analysis_type = "URL-only BERT Analysis"
            fetch_status = f"⚠️ Could not fetch HTML content: {status}"
    else:
        # Process as regular text
        probs, tokens, has_attention, attention_weights = predict_bert_single_text(text, "text")
        analysis_type = "BERT Text Analysis"
        fetch_status = ""
    
    # Create detailed analysis
    predicted_class = "phishing" if probs[1] > probs[0] else "benign"
    confidence = max(probs)
    
    detailed_analysis = f"""
<div style="font-family: Arial, sans-serif; max-width: 800px; margin: 0 auto; background: #1e1e1e; padding: 20px; border-radius: 15px;">
<div style="background: linear-gradient(135deg, {'#8b0000' if predicted_class == 'phishing' else '#006400'} 0%, {'#dc143c' if predicted_class == 'phishing' else '#228b22'} 100%); padding: 25px; border-radius: 20px; color: white; text-align: center; margin-bottom: 20px; box-shadow: 0 8px 32px rgba(0,0,0,0.5); border: 2px solid {'#ff4444' if predicted_class == 'phishing' else '#44ff44'};">
    <h2 style="margin: 0 0 10px 0; font-size: 28px; color: white;">🔍 {analysis_type}</h2>
    <div style="font-size: 36px; font-weight: bold; margin: 10px 0; color: white;">
        {predicted_class.upper()}
    </div>
    <div style="font-size: 18px; color: #f0f0f0;">
        Confidence: {confidence:.1%}
    </div>
    <div style="margin-top: 15px; font-size: 14px; color: #e0e0e0;">
        {'This appears to be a phishing attempt!' if predicted_class == 'phishing' else '✅ This appears to be legitimate content.'}
    </div>
</div>
"""

    if fetch_status:
        detailed_analysis += f"""
<div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
    <strong>Fetch Status:</strong> {fetch_status}
</div>
"""
    
    if has_attention and attention_weights is not None:
        attention_scores = attention_weights.squeeze(0).tolist() if attention_weights.dim() > 1 else attention_weights.tolist()
        
        token_analysis = []
        for i, (token, score) in enumerate(zip(tokens, attention_scores)):
            # More lenient filtering - include more tokens for text analysis
            if token not in ['[CLS]', '[SEP]', '[PAD]', '<s>', '</s>'] and len(token.strip()) > 0 and score > 0.005:  
                clean_token = token.replace('▁', '').replace('Ġ', '').strip()  # Handle different tokenizer prefixes
                if clean_token:  # Only add if token has content after cleaning
                    token_analysis.append({
                        'token': clean_token,
                        'importance': score,
                        'position': i
                    })
        
        # Sort by importance
        token_analysis.sort(key=lambda x: x['importance'], reverse=True)
        
        detailed_analysis += f"""
## Top important tokens:
<div style="background: #2d2d2d; padding: 15px; border-radius: 10px; margin: 15px 0; border-left: 4px solid #4caf50; color: #e0e0e0;">
    <strong>Analysis Info:</strong> Found {len(token_analysis)} important tokens out of {len(tokens)} total tokens
</div>
<div style="font-family: Arial, sans-serif;">
"""
        
        for i, token_info in enumerate(token_analysis[:10]):  # Top 10 tokens
            bar_width = int(token_info['importance'] * 100)
            color = "#ff4444" if predicted_class == "phishing" else "#44ff44"
            
            detailed_analysis += f"""
<div style="margin: 8px 0; display: flex; align-items: center; background: #2d2d2d; padding: 8px; border-radius: 8px; border-left: 4px solid {color};">
    <div style="width: 30px; text-align: right; margin-right: 10px; font-weight: bold; color: #ffffff;">
        {i+1}.
    </div>
    <div style="width: 120px; margin-right: 10px; font-weight: bold; color: #e0e0e0; text-align: right;">
        {token_info['token']}
    </div>
    <div style="width: 300px; background-color: #404040; border-radius: 10px; overflow: hidden; margin-right: 10px; border: 1px solid #555;">
        <div style="width: {bar_width}%; background-color: {color}; height: 20px; border-radius: 10px; transition: width 0.3s ease;"></div>
    </div>
    <div style="color: #cccccc; font-size: 12px; font-weight: bold;">
        {token_info['importance']:.1%}
    </div>
</div>
"""
        
        detailed_analysis += "</div>\n"
        
        detailed_analysis += f"""
## Detailed analysis:
<div style="font-family: Arial, sans-serif; background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
    <h3 style="margin: 0 0 15px 0; color: white;">Statistical Overview</h3>
    <div style="display: grid; grid-template-columns: repeat(2, 1fr); gap: 15px;">
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Total tokens</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in token_analysis if t['importance'] > 0.05])}</div>
            <div style="font-size: 14px; color: #e0e0e0;">High impact tokens (>5%)</div>
        </div>
    </div>
</div>
<div style="font-family: Arial, sans-serif; margin: 15px 0; background: #2d2d2d; padding: 20px; border-radius: 15px; border: 1px solid #555;">
    <h3 style="color: #ffffff; margin-bottom: 15px;"> Prediction Confidence</h3>
    <div style="display: flex; justify-content: space-between; margin-bottom: 10px;">
        <span style="font-weight: bold; color: #ff4444;">Phishing</span>
        <span style="font-weight: bold; color: #44ff44;">Benign</span>
    </div>
    <div style="width: 100%; background-color: #404040; border-radius: 25px; overflow: hidden; height: 30px; border: 1px solid #666;">
        <div style="width: {probs[1]*100:.1f}%; background: linear-gradient(90deg, #ff4444 0%, #ff6666 100%); height: 100%; display: flex; align-items: center; justify-content: center; color: white; font-weight: bold; font-size: 14px;">
            {probs[1]:.1%}
        </div>
    </div>
    <div style="margin-top: 10px; text-align: center; color: #cccccc; font-size: 14px;">
        Benign: {probs[0]:.1%}
    </div>
</div>
"""
    else:
        # Fallback analysis without attention weights
        detailed_analysis += f"""
<div style="background: linear-gradient(135deg, #1a237e 0%, #3949ab 100%); padding: 20px; border-radius: 15px; color: white; margin: 15px 0; border: 2px solid #3f51b5;">
    <h3 style="margin: 0 0 15px 0; color: white;">Basic Analysis</h3>
    <div style="display: grid; grid-template-columns: repeat(3, 1fr); gap: 15px;">
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[1]:.1%}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Phishing</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{probs[0]:.1%}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Benign</div>
        </div>
        <div style="background: rgba(255,255,255,0.1); padding: 15px; border-radius: 10px; text-align: center; border: 1px solid rgba(255,255,255,0.2);">
            <div style="font-size: 24px; font-weight: bold; color: white;">{len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])}</div>
            <div style="font-size: 14px; color: #e0e0e0;">Tokens</div>
        </div>
    </div>
</div>
<div style="background: #2d2d2d; padding: 20px; border-radius: 15px; margin: 15px 0; border: 1px solid #555;">
    <h3 style="color: #ffffff; margin: 0 0 15px 0;">🔤 Tokens in text:</h3>
    <div style="display: flex; flex-wrap: wrap; gap: 8px;">""" + ''.join([f'<span style="background: #404040; color: #64b5f6; padding: 4px 8px; border-radius: 15px; font-size: 12px; border: 1px solid #666;">{token.replace("▁", "")}</span>' for token in tokens if token not in ['[CLS]', '[SEP]', '[PAD]']]) + f"""</div>
    <div style="margin-top: 15px; padding: 10px; background: #3d2914; border-radius: 8px; border-left: 4px solid #ff9800;">
        <strong style="color: #ffcc02;">Debug info:</strong> <span style="color: #e0e0e0;">Found {len(tokens)} total tokens, {len([t for t in tokens if t not in ['[CLS]', '[SEP]', '[PAD]']])} content tokens</span>
    </div>
</div>
<div style="background: #3d2914; padding: 15px; border-radius: 10px; border-left: 4px solid #ff9800; margin: 15px 0;">
    <p style="margin: 0; color: #ffcc02; font-size: 14px;">
        <strong>Note:</strong> Detailed attention weights analysis is not available for the current model.
    </p>
</div>
"""

    detailed_analysis += "</div>"
    
    # Build label->prob mapping for Gradio Label output
    if len(LABELS) == len(probs):
        prediction_result = {LABELS[i]: float(probs[i]) for i in range(len(LABELS))}
    else:
        prediction_result = {f"class_{i}": float(p) for i, p in enumerate(probs)}
    
    return prediction_result, detailed_analysis



# --------- Gradio UI ----------
deberta_interface = gr.Interface(
    fn=predict_fn,
    inputs=gr.Textbox(label="URL or text", placeholder="Example: http://suspicious-site.example or paste any text"),
    outputs=[
        gr.Label(label="Prediction result"),
        gr.Markdown(label="Detailed token analysis")
    ],
    title="Phishing Detector (DeBERTa + LSTM)",
    description="""
    Enter a URL or text for analysis.
    **Features:**
    - **URL Analysis**: For URLs, the system will fetch HTML content and combine both URL and content analysis
    - **Combined Prediction**: Uses weighted combination of URL structure and webpage content analysis  
    - **Visual Analysis**: Predict phishing/benign probability with visual charts
    - **Token Importance**: Display the most important tokens in classification
    - **Detailed Insights**: Comprehensive analysis of the impact of each token
    - **Dark Theme**: Beautiful interface with colorful charts optimized for dark themes
    
    **How it works for URLs:**
    1. Analyze the URL structure itself
    2. Fetch the webpage HTML content
    3. Analyze the webpage content
    4. Combine both results for final prediction (30% URL + 70% content)
    """,
    examples=[
        ["http://rendmoiunserviceeee.com"],
        ["https://www.google.com"],
        ["Dear customer, your account has been suspended. Click here to verify your identity immediately."],
        ["https://mail-secure-login-verify.example/path?token=suspicious"],
        ["http://paypaI-security-update.net/login"],
        ["Your package has been delivered successfully. Thank you for using our service."],
        ["https://github.com/user/repo"]
    ],
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        background-color: #1e1e1e !important;
        color: #ffffff !important;
    }
    .dark .gradio-container {
        background-color: #1e1e1e !important;
    }
    /* Dark theme for all components */
    .block {
        background-color: #2d2d2d !important;
        border: 1px solid #444 !important;
    }
    .gradio-textbox {
        background-color: #3d3d3d !important;
        color: #ffffff !important;
        border: 1px solid #666 !important;
    }
    .gradio-button {
        background-color: #4a4a4a !important;
        color: #ffffff !important;
        border: 1px solid #666 !important;
    }
    .gradio-button:hover {
        background-color: #5a5a5a !important;
    }
    """
)

bert_interface = gr.Interface(
    fn=predict_bert_interface_fn,
    inputs=gr.Textbox(label="URL or text", placeholder="Example: http://suspicious-site.example or paste any text"),
    outputs=[
        gr.Label(label="Prediction result"),
        gr.Markdown(label="Detailed token analysis")
    ],
    title="Phishing Detector (BERT)",
    description="""
    Enter a URL or text for analysis using the BERT model.
    **Features:**
    - **URL Analysis**: For URLs, the system will fetch HTML content and combine both URL and content analysis
    - **Combined Prediction**: Uses weighted combination of URL structure and webpage content analysis  
    - **Visual Analysis**: Predict phishing/benign probability with visual charts
    - **Token Importance**: Display the most important tokens in classification using attention weights
    - **Detailed Insights**: Comprehensive analysis of the impact of each token
    - **Dark Theme**: Beautiful interface with colorful charts optimized for dark themes
    
    **How it works for URLs:**
    1. Analyze the URL structure itself
    2. Fetch the webpage HTML content
    3. Analyze the webpage content
    4. Combine both results for final prediction (30% URL + 70% content)
    """,
    examples=[
        ["http://rendmoiunserviceeee.com"],
        ["https://www.google.com"],
        ["Dear customer, your account has been suspended. Click here to verify your identity immediately."],
        ["https://mail-secure-login-verify.example/path?token=suspicious"],
        ["http://paypaI-security-update.net/login"],
        ["Your package has been delivered successfully. Thank you for using our service."],
        ["https://github.com/user/repo"]
    ],
    theme=gr.themes.Soft(),
    css="""
    .gradio-container {
        font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
        background-color: #1e1e1e !important;
        color: #ffffff !important;
    }
    .dark .gradio-container {
        background-color: #1e1e1e !important;
    }
    /* Dark theme for all components */
    .block {
        background-color: #2d2d2d !important;
        border: 1px solid #444 !important;
    }
    .gradio-textbox {
        background-color: #3d3d3d !important;
        color: #ffffff !important;
        border: 1px solid #666 !important;
    }
    .gradio-button {
        background-color: #4a4a4a !important;
        color: #ffffff !important;
        border: 1px solid #666 !important;
    }
    .gradio-button:hover {
        background-color: #5a5a5a !important;
    }
    """
)



demo = gr.TabbedInterface(
    [deberta_interface, bert_interface],
    ["DeBERTa + LSTM", "BERT"]
)

if __name__ == "__main__":
    demo.launch()