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."
)

model_path = "./URLGuardian" 
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 Safety of the url"):
    if url:
        # Run the classifier on the input URL
        result = classifier(url)[0]
        label = result["label"]
        score = result["score"]

        # Display result based on the label
        # Adjust the label checking logic based on the model's documentation.
        if "safe" in label.lower():
            st.success(
                f"The URL '{url}' is considered safe with a confidence of {score * 100:.2f}%."
            )
        else:
            st.error(
                f"The URL '{url}' is considered suspicious with a confidence of {score * 100:.2f}%."
            )
        # Optionally, you can display the full result for debugging purposes:
        st.write("Full classification output:", result)
    else:
        st.error("Please enter a URL.")