import pandas as pd
import gradio as gr
import joblib

pipeline = joblib.load("pipeline.joblib")

def predict(
    Have_IP, 
    Have_At, 
    URL_Length, 
    URL_Depth, 
    Redirection, 
    https_Domain, 
    TinyURL, 
    Prefix_or_Suffix, 
    DNS_Record,
    Domain_Age,
    Domain_End,
    iFrame,
    Mouse_Over,
    Right_Click,
    Web_Forwards
    ):

    sample = {
        "Have_IP": Have_IP,
        "Have_At": Have_At,
        "URL_Length": URL_Length,
        "URL_Depth": URL_Depth,
        "Redirection": Redirection,
        "https_Domain": https_Domain,
        "TinyURL": TinyURL,
        "Prefix/Suffix": Prefix_or_Suffix,
        "DNS_Record": DNS_Record,
        "Domain_Age": Domain_Age,
        "Domain_End": Domain_End,
        "iFrame": iFrame,
        "Mouse_Over": Mouse_Over,
        "Right_Click": Right_Click,
        "Web_Forwards": Web_Forwards
    }

    sample_df = pd.DataFrame([sample])

    y_pred = pipeline.predict(sample_df)
    y_pred = int(y_pred[0])

    return y_pred

with gr.Blocks() as demo:
    gr.Markdown("# URL Classification Demo")
    Have_IP = gr.Checkbox(label="Have_IP")
    Have_At = gr.Checkbox(label="Have_At")
    URL_Length = gr.Checkbox(label="URL_Length")
    URL_Depth = gr.Number(label="URL_Depth", value=0)
    Redirection = gr.Checkbox(label="Redirection")
    https_Domain = gr.Checkbox(label="https_Domain")
    TinyURL = gr.Checkbox(label="TinyURL")
    Prefix_or_Suffix = gr.Checkbox(label="Prefix/Suffix")
    DNS_Record = gr.Checkbox(label="DNS_Record")
    Domain_Age = gr.Checkbox(label="Domain_Age")
    Domain_End = gr.Checkbox(label="Domain_End")
    iFrame = gr.Checkbox(label="iFrame")
    Mouse_Over = gr.Checkbox(label="Mouse_Over")
    Right_Click = gr.Checkbox(label="Right_Click")
    Web_Forwards = gr.Checkbox(label="Web_Forwards")
    
    output = gr.Label(label="Prediction Output")

    inputs = [
        Have_IP, Have_At, URL_Length, URL_Depth, Redirection, https_Domain, TinyURL, 
        Prefix_or_Suffix, DNS_Record, Domain_Age, Domain_End, iFrame, 
        Mouse_Over, Right_Click, Web_Forwards
    ]
    outputs = output

    predict_btn = gr.Button("Predict", variant="primary")
    predict_btn.click(predict, inputs=inputs, outputs=outputs)

if __name__ == "__main__":
    demo.launch()