import gradio as gr
import joblib
import pandas as pd

# Load trained model
model = joblib.load("cybersecurity_model.joblib")

# List of features your model expects (update as per your actual features)
FEATURE_NAMES = [
    'A_frequency','NS_frequency','CNAME_frequency','SOA_frequency','NULL_frequency',
    'PTR_frequency','HINFO_frequency','MX_frequency','TXT_frequency','AAAA_frequency',
    'SRV_frequency','OPT_frequency','rr_type','rr_count','rr_name_entropy','rr_name_length',
    'distinct_ns','distinct_ip','unique_country','unique_asn','distinct_domains','reverse_dns',
    'a_records','unique_ttl','ttl_mean','ttl_variance','FQDN_count','subdomain_length','upper',
    'lower','numeric','entropy','special','labels','labels_max','labels_average','longest_word',
    'sld','len','subdomain'
]

# Inference function
def predict_from_csv(file):
    try:
        df = pd.read_csv(file)
        
        # Ensure required columns are present
        missing_cols = set(FEATURE_NAMES) - set(df.columns)
        if missing_cols:
            return f"❌ Missing columns in CSV: {', '.join(missing_cols)}"

        # Select and predict
        X = df[FEATURE_NAMES]
        preds = model.predict(X)
        
        # Format output nicely
        result = pd.DataFrame({"Prediction": preds})
        return result
    except Exception as e:
        return f"❌ Error reading file: {str(e)}"

# Gradio UI
interface = gr.Interface(
    fn=predict_from_csv,
    inputs=gr.File(label="📁 Upload CSV File with Network Traffic Features"),
    outputs=gr.Dataframe(label="📊 Predictions"),
    title="🚨 Cybersecurity Attack Detector",
    description=(
        "AI-powered model to detect attacks from DNS/network traffic data. "
        "Upload a CSV file with preprocessed features to get predictions."
    ),
    allow_flagging="never"
)

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