import gradio as gr
import joblib
import numpy as np
import json
from huggingface_hub import hf_hub_download

# ========================
# DOWNLOAD MODEL FILES
# ========================
repo_id = "Daksh159/soc-ai-detection"

lgb_path = hf_hub_download(repo_id, "lgb_model.pkl")
scaler_path = hf_hub_download(repo_id, "scaler.pkl")
le_path = hf_hub_download(repo_id, "label_encoder.pkl")
feature_path = hf_hub_download(repo_id, "feature_names.json")

# Load models
lgb_model = joblib.load(lgb_path)
scaler = joblib.load(scaler_path)
le = joblib.load(le_path)

with open(feature_path) as f:
    feature_names = json.load(f)

# ========================
# MITRE KEYWORDS (ROBUST)
# ========================
mitre_mapping = {
    "sql injection": ("T1190", "Exploit Public-Facing Application"),
    "ddos": ("T1498", "Network Denial of Service"),
    "portscan": ("T1046", "Network Service Scanning"),
    "brute force": ("T1110", "Brute Force")
}

# ========================
# PREDICT FUNCTION
# ========================
def predict(input_text):
    try:
        # Parse JSON
        data = json.loads(input_text)

        # Convert to ordered feature list
        values = [data.get(f, 0) for f in feature_names]

        x = np.array(values).reshape(1, -1)
        x = scaler.transform(x)

        # Prediction
        pred = lgb_model.predict(x)[0]
        label = le.inverse_transform([pred])[0]

        # CLEAN LABEL
        label = label.replace('\ufffd', '')
        label = " ".join(label.split()).strip()

        # Confidence
        proba = lgb_model.predict_proba(x)[0]
        confidence = float(np.max(proba))
        confidence = min(confidence, 0.99)

        # ========================
        # MITRE MATCHING
        # ========================
        normalized_label = label.lower()
        mitre_id, mitre_name = "Unknown", "Unknown"

        for key in mitre_mapping:
            if key in normalized_label:
                mitre_id, mitre_name = mitre_mapping[key]
                break

        # Final Output
        return json.dumps({
            "prediction": label,
            "confidence": round(confidence, 3),
            "mitre": f"{mitre_id} - {mitre_name}"
        }, indent=4)

    except Exception as e:
        return f"Error: {str(e)}"


# ========================
# UI
# ========================
iface = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(
        label="Paste JSON logs",
        placeholder='{"Flow Duration": 10000, "ACK Flag Count": 2}'
    ),
    outputs="text",
    title="SOC AI Detection System",
    description="Paste structured JSON logs from Elastic"
)

iface.launch()