import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model = AutoModelForSequenceClassification.from_pretrained("duclo90/PhishingClassifier")
tokenizer = AutoTokenizer.from_pretrained("duclo90/PhishingClassifier")
model.eval()

label_map = {0: "Safe", 1: "Phishing"}

def classify(text):
    if not text.strip():
        return "⚠️ Please enter some text to analyze", None
    
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
        probs = torch.nn.functional.softmax(logits, dim=-1)
        pred = torch.argmax(logits, dim=-1).item()
        confidence = probs[0][pred].item() * 100
    
    result = label_map[pred]
    
    if result == "Safe":
        status = f"✅ **SAFE** - This content appears legitimate"
        color_indicator = "🟢"
    else:
        status = f"🚨 **PHISHING DETECTED** - This content may be malicious"
        color_indicator = "🔴"
    
    detailed_result = f"""
{color_indicator} **Result:** {result}
📊 **Confidence:** {confidence:.2f}%

{status}
    """
    
    return detailed_result.strip(), confidence

# Custom CSS for a modern cybersecurity aesthetic
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap');

* {
    font-family: 'Inter', sans-serif !important;
}

.gradio-container {
    background: linear-gradient(135deg, #0f0f23 0%, #1a1a2e 50%, #16213e 100%) !important;
    color: #e0e0e0 !important;
}

#component-0 {
    max-width: 900px !important;
    margin: 0 auto !important;
    padding: 2rem !important;
}

.contain {
    background: rgba(255, 255, 255, 0.03) !important;
    backdrop-filter: blur(10px) !important;
    border: 1px solid rgba(255, 255, 255, 0.1) !important;
    border-radius: 16px !important;
    padding: 2rem !important;
    box-shadow: 0 8px 32px rgba(0, 0, 0, 0.3) !important;
}

.input-text textarea {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 2px solid rgba(100, 200, 255, 0.3) !important;
    border-radius: 12px !important;
    color: #e0e0e0 !important;
    font-size: 16px !important;
    padding: 1rem !important;
    transition: all 0.3s ease !important;
}

.input-text textarea:focus {
    border-color: rgba(100, 200, 255, 0.6) !important;
    box-shadow: 0 0 20px rgba(100, 200, 255, 0.2) !important;
    outline: none !important;
}

button.primary {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    border-radius: 12px !important;
    color: white !important;
    font-weight: 600 !important;
    padding: 0.75rem 2rem !important;
    font-size: 16px !important;
    transition: all 0.3s ease !important;
    box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
}

button.primary:hover {
    transform: translateY(-2px) !important;
    box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
}

.output-text {
    background: rgba(255, 255, 255, 0.05) !important;
    border: 2px solid rgba(100, 200, 255, 0.2) !important;
    border-radius: 12px !important;
    padding: 1.5rem !important;
    color: #e0e0e0 !important;
    font-size: 16px !important;
    line-height: 1.8 !important;
}

h1 {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    font-weight: 700 !important;
    font-size: 2.5rem !important;
    margin-bottom: 0.5rem !important;
    text-align: center !important;
}

.description {
    color: #b0b0b0 !important;
    text-align: center !important;
    font-size: 1.1rem !important;
    margin-bottom: 2rem !important;
}

.footer {
    text-align: center !important;
    margin-top: 2rem !important;
    padding-top: 1.5rem !important;
    border-top: 1px solid rgba(255, 255, 255, 0.1) !important;
    color: #808080 !important;
    font-size: 0.9rem !important;
}

.progress {
    background: rgba(100, 200, 255, 0.2) !important;
    border-radius: 8px !important;
}

.progress-bar {
    background: linear-gradient(90deg, #667eea 0%, #764ba2 100%) !important;
}
"""

# Create the interface with enhanced design
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🛡️ Phishing Content Detector
        ### AI-Powered Security Analysis
        """
    )
    
    gr.Markdown(
        """
        <p class="description">
        Protect yourself from phishing attacks. Paste suspicious emails, messages, or text below for instant AI analysis.
        </p>
        """,
        elem_classes="description"
    )
    
    with gr.Row():
        with gr.Column(scale=1):
            input_text = gr.Textbox(
                label="📝 Content to Analyze",
                placeholder="Paste suspicious email, message, or text here...\n\nExample: 'Your account has been locked. Click here immediately to verify your identity and avoid suspension.'",
                lines=8,
                elem_classes="input-text"
            )
            
            analyze_btn = gr.Button("🔍 Analyze Content", variant="primary", size="lg")
            
            gr.Markdown(
                """
                <div class="footer">
                <strong>💡 Tips:</strong> Look for urgent language, suspicious links, requests for personal information, or grammar errors.
                <br>
                <em>Powered by AI • Model: duclo90/PhishingClassifier</em>
                </div>
                """
            )
    
    with gr.Row():
        with gr.Column(scale=1):
            output_text = gr.Textbox(
                label="🎯 Analysis Result",
                lines=6,
                elem_classes="output-text"
            )
            
            confidence_slider = gr.Slider(
                label="Confidence Level",
                minimum=0,
                maximum=100,
                value=0,
                interactive=False,
                elem_classes="progress"
            )
    
    # Examples section
    gr.Examples(
        examples=[
            ["Congratulations! You've won $1,000,000! Click here now to claim your prize before it expires!"],
            ["Hi team, the quarterly meeting is scheduled for next Tuesday at 2 PM in Conference Room B."],
            ["URGENT: Your account will be suspended. Verify your identity immediately by clicking this link."],
            ["Your package delivery failed. Update your address at: legitimate-shipping-company.com"],
        ],
        inputs=input_text,
        label="📋 Try These Examples"
    )
    
    analyze_btn.click(
        fn=classify,
        inputs=input_text,
        outputs=[output_text, confidence_slider]
    )
    
    input_text.submit(
        fn=classify,
        inputs=input_text,
        outputs=[output_text, confidence_slider]
    )

demo.launch()