import gradio as gr
from transformers import BertForSequenceClassification, BertTokenizer
import torch
import torch.nn.functional as F

# Load model and tokenizer
model_name = 'ealvaradob/bert-finetuned-phishing'
model = BertForSequenceClassification.from_pretrained(model_name)
tokenizer = BertTokenizer.from_pretrained(model_name)
model.eval()

# Prediction function with confidence
def phishCheck(email_text):
    inputs = tokenizer(email_text, return_tensors="pt", truncation=True, padding='max_length', max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = F.softmax(logits, dim=-1).squeeze()

    predicted_class = torch.argmax(probs).item()
    confidence = probs[predicted_class].item()
    label = "Phishing" if predicted_class == 1 else "Legitimate"
    confidence_pct = round(confidence * 100, 2)

    return label, confidence_pct

# Format HTML table
def format_history_as_html(history):
    table = "<table style='width: 100%; border-collapse: collapse;'>"
    table += "<tr><th style='border: 1px solid black; padding: 5px;'>Email Text</th><th style='border: 1px solid black; padding: 5px;'>Result</th><th style='border: 1px solid black; padding: 5px;'>Confidence</th></tr>"
    for email, label, confidence in history:
        color = "green" if label == "Legitimate" else "red"
        table += f"<tr><td style='border: 1px solid black; padding: 5px;'>{email}</td>"
        table += f"<td style='border: 1px solid black; padding: 5px; color: white; background-color: {color}; text-align: center;'>{label}</td>"
        table += f"<td style='border: 1px solid black; padding: 5px; text-align: center;'>{confidence:.2f}%</td></tr>"
    table += "</table>"
    return table

# Main function
def process_input(email_body, history):
    if not email_body.strip():  # check for empty or whitespace input
        styled_result = """
        <div style="color: black; background-color: yellow; padding: 10px; border-radius: 5px; text-align: center;">
            No text detected
        </div>
        """

        recommendation = """
        <div style='background-color: #ffffcc; padding: 10px; border-left: 5px solid #cccc00; color: black;'>
            ⚠️ <strong style='color: black;'>Recommendation:</strong> Please input text to analyze.
        </div>
        """

        return styled_result, recommendation, history, format_history_as_html(history)

    # If there is valid input, proceed as normal
    label, confidence_pct = phishCheck(email_body)

    # Style result box
    color = "green" if label == "Legitimate" else "red"
    styled_result = f'<div style="color: white; background-color: {color}; padding: 10px; border-radius: 5px; text-align: center;">{label} ({confidence_pct}%)</div>'
   
    # Recommendation logic
    if label == "Phishing":
        recommendation = """
        <div style='background-color: #ffe6e6; padding: 10px; border-left: 5px solid red; color: black;'>
        <span style='color: black;'>⚠️ <strong style='color: black;'>Recommendation:</strong> This message appears to be phishing. Do not click any links, download attachments, or provide sensitive information. Report this message to your IT or security team, if applicable.</span>
        </div>
        """
    else:
        recommendation = """
        <div style='background-color: #e6ffe6; padding: 10px; border-left: 5px solid green; color: black;'>
        <span style='color: black;'>✅ <strong style='color: black;'>Recommendation:</strong> This message appears legitimate. Still, use caution when interacting with links, attachments, or requests for sensitive information.</span>
        </div>
        """

    history.append((email_body, label, confidence_pct))
    history_html = format_history_as_html(history)

    return styled_result, recommendation, history, history_html


# Interface with loading
with gr.Blocks() as demo:
    gr.Markdown("## 🕵️‍♂️ Phishing Message Detector with Confidence Score and History")
    gr.Markdown("""
    ###  Project Description
    This is an experimental tool designed by TechHorizon Consulting to help users quickly assess the likelihood that an email or text message may be a phishing attempt. By analyzing language patterns and known phishing indicators, this tool provides an AI-generated risk assessment to support safer decision-making.
    To use this tool, copy and paste the message body into the text area below, then click the “Scan for Phishing” button. Please wait several seconds for processing, then review the results and recommended action. 
 
    """)

    gr.Markdown("""
    ### ⚠️ Disclaimer

    This tool is provided for **informational and educational purposes only**. It is **not a substitute** for professional cybersecurity services and should **not be relied upon** as the sole method for detecting phishing, malware, or other malicious content.

    **Do not submit** any personally identifiable information (PII), passwords, sensitive financial data, or protected health information (PHI). All input is processed by a **public AI model** and may be stored temporarily for performance or logging purposes.

    By using this tool, you agree that:

    - You are solely responsible for the content you submit.
    - The tool does not guarantee accuracy or security outcomes.
    - The creators assume no liability for any damages resulting from the use or misuse of this tool or its outputs.

    For any suspicious messages, **always follow your organization’s cybersecurity protocols** and consult a qualified IT or security professional.
    """)

    input_box = gr.Textbox(label="Paste Message Here:")  ##laceholder="Paste your text here")
    detect_btn = gr.Button("Scan For Phishing")

    output_box = gr.HTML(
    value="<div style='text-align: center; padding: 10px; color: #666;'>Waiting for input...</div>",
    label="Detection Result"
)
    
    recommendation_box = gr.HTML(
        value="<div style='text-align: center; padding: 10px; color: #666;'>No recommendation yet.</div>",
        label="Recommendation"
    )

    history_box = gr.HTML(label="Input History (Email + Result + Confidence)")

    history_state = gr.State([])

    detect_btn.click(
        fn=process_input,
        inputs=[input_box, history_state],
        outputs=[output_box, recommendation_box, history_state, history_box],
        show_progress=True
    )


demo.launch()