# STEP 1: Install Dependencies
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F
import gradio as gr


# STEP 2: Clone the model
!git lfs install
!git clone https://huggingface.co/ealvaradob/bert-finetuned-phishing

# STEP 3: Load Model & Tokenizer
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

model_path = "./bert-finetuned-phishing"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForSequenceClassification.from_pretrained(model_path)
model.eval()

# STEP 4: Prediction Function
def predict_phishing_url(url):
    inputs = tokenizer(url, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        probs = F.softmax(outputs.logits, dim=1)
        predicted_class = torch.argmax(probs).item()
        confidence = probs[0][predicted_class].item()
        
    label = "🚨 Phishing Detected!" if predicted_class == 1 else "✅ Safe URL"
    return f"{label} (Confidence: {confidence:.2f})"

# STEP 5: Custom CSS and Gradio Interface
custom_css = """
body {
    background-image: url('https://th.bing.com/th/id/OIP.smzBbDZkYqUVkQQ6yFWmHQAAAA?w=251&h=180&c=7&r=0&o=5&dpr=2&pid=1.7');
    background-size: cover;
    background-attachment: fixed;
    background-position: center;
    color: white;
    display: flex;
    justify-content: center;
    align-items: center;
    height: 100vh;
}
h1, h2, h3 {
    color: #00ffcc;
    text-shadow: 0 0 5px #00ffcc;
    text-align: center;
}
.gradio-container {
    background: rgba(0,0,0,0.8);
    border: 2px solid #00ffcc;
    border-radius: 20px;
    padding: 30px;
    max-width: 600px;
    margin: auto;
}
input, textarea {
    background: rgba(0,0,0,0.5);
    color: #00ffcc;
    border: 1px solid #00ffcc;
}
button {
    background: #00ffcc;
    color: black;
    border-radius: 10px;
    padding: 10px 20px;
    font-weight: bold;
    margin-top: 10px;
}
"""

import gradio as gr

with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_id="center-col"):
        gr.Markdown("## 🕵️‍♂️ Phishing URL Detector - BERT + Hacker Style Web App")
        gr.Markdown("### Enter a URL to check if it's phishing or safe.")
        url_input = gr.Textbox(label="Enter URL", placeholder="https://example.com/secure-login")
        output = gr.Textbox(label="Result")
        check_btn = gr.Button("Check URL")
        check_btn.click(fn=predict_phishing_url, inputs=url_input, outputs=output)

demo.launch()