import streamlit as st
import torch
from PIL import Image
from transformers import pipeline, AutoModelForSequenceClassification, AutoTokenizer

# Initialize the image-to-text pipeline and models
@st.cache(allow_output_mutation=True)
def load_models():
    image_pipeline = pipeline("image-to-text", model="microsoft/trocr-large-printed")
    phishing_model = AutoModelForSequenceClassification.from_pretrained("kithangw/phishing_link_detection", num_labels=2)
    phishing_tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-2_H-128_A-2")
    return image_pipeline, phishing_model, phishing_tokenizer

image_pipeline, phishing_model, phishing_tokenizer = load_models()

# Define the phishing check function
def check_phishing(url_for_recognize):
    link_token = phishing_tokenizer(url_for_recognize, max_length=512, padding=True, truncation=True, return_tensors='pt')

    with torch.no_grad():  # Disable gradient calculation for inference
        output = phishing_model(**link_token)

    probabilities = torch.nn.functional.softmax(output.logits, dim=-1)
    predicted_class = torch.argmax(probabilities, dim=-1).item()
    predicted_prob = probabilities[0, predicted_class].item()

    labels = ['Not Phishing', 'Phishing']
    prediction_label = labels[predicted_class]
    sentence = f"The URL '{url_for_recognize}' is classified as '{prediction_label}' with a probability of {predicted_prob:.2f}."
    return sentence

# Streamlit interface
st.title("Phishing URL Detection from Image")
uploaded_image = st.file_uploader("Upload an image of the URL", type=["png", "jpg", "jpeg"])

if uploaded_image is not None:
    image = Image.open(uploaded_image)
    st.image(image, caption='Uploaded URL Image', use_column_width=True)
    
    # Convert image to URL text
    url_for_recognize = image_pipeline(uploaded_image)[0]['generated_text'].replace(" ", "").lower()
    st.write("Recognized URL:")
    # Use a text input to let the user verify and possibly edit the recognized URL
    verified_url = st.text_input("Verify or edit the recognized URL if necessary:", value=url_for_recognize)

    if st.button('Detect Phishing'):
        if verified_url:
            result = check_phishing(verified_url)
            st.write(result)
        else:
            st.write("Please enter a URL to check for phishing.")