import gradio as gr
import pandas as pd
import joblib

# Load mô hình và các công cụ đã lưu
model = joblib.load("xgboost_model.pkl")
label_encoder = joblib.load("label_encoder.pkl")
scaler = joblib.load("scaler.pkl")

# Hàm dự đoán từ file CSV
def predict_from_csv(file):
    try:
        df = pd.read_csv(file.name)
        original_rows = len(df)

        # Tiền xử lý giống như khi train
        scaled_data = scaler.transform(df)
        preds = model.predict(scaled_data)

        # Giải mã nhãn nếu cần
        labels = label_encoder.inverse_transform(preds)

        # Gắn kết quả vào DataFrame
        df["Dự đoán"] = labels
        return df

    except Exception as e:
        return f"Lỗi: {str(e)}"

# Giao diện Gradio
iface = gr.Interface(
    fn=predict_from_csv,
    inputs=gr.File(label="Tải file CSV để dự đoán"),
    outputs=gr.Dataframe(label="Kết quả dự đoán"),
    title="🧠 Dự đoán tấn công mạng bằng mô hình XGBoost",
    description="Tải file .csv (đã qua xử lý đầu vào giống lúc huấn luyện) để hệ thống phân loại và dự đoán loại tấn công."
)

iface.launch()