import streamlit as st
from transformers import pipeline

@st.cache_resource
def load_classifier(model_path: str):
    # Loads the URLGuardian classifier from the Hugging Face Hub.
    return pipeline("text-classification", model=model_path)

# App Title and description
st.title("URL Typosquatting Detection with URLGuardian")
st.markdown(
    "This app uses the **URLGuardian** classifier by Anvilogic from the Hugging Face Hub to detect potential typosquatting. "
    "Enter a legitimate URL and a potentially typosquatted URL to see the classifier's prediction."
)

# Load the classifier model from Hugging Face
model_path = "Anvilogic/URLGuardian"  # Model repository on Hugging Face
classifier = load_classifier(model_path)

# URL inputs
url = st.text_input("Enter the URL:", value="https://example.com")

# Typosquatting detection on button click
if st.button("Check Typosquatting"):
    if legit_url and suspect_url:
        
        result = classifier(url)[0]
        label = result['label']
        score = result['score']
        
        # Adjust the label names as per the model's documentation.
        # This example assumes the label for a typosquatted URL might include "typo".
        if "typo" in label.lower():
            st.success(
                f"The model predicts that '{suspect_url}' is likely a typosquatted version of '{legit_url}' "
                f"with a confidence of {score * 100:.2f}%."
            )
        else:
            st.warning(
                f"The model predicts that '{suspect_url}' is NOT likely a typosquatted version of '{legit_url}' "
                f"with a confidence of {score * 100:.2f}%."
            )
    else:
        st.error("Please enter both a legitimate URL and a potentially typosquatted URL.")