import gradio as gr
import pandas as pd
import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split

# --- BAGIAN 1: MEMBUAT DATASET LALU LINTAS JARINGAN (SIMULASI) ---
# Di dunia nyata, ini diganti dengan load file seperti NSL-KDD.csv
def create_dummy_dataset():
    # Kita buat 1000 baris data dummy
    n_samples = 1000
    np.random.seed(42)
    
    # Fitur 1: Durasi Koneksi (detik)
    # Normal: cepat (0-5 detik), Serangan (DoS): lama (10-100 detik)
    duration = np.concatenate([np.random.randint(0, 5, 900), np.random.randint(20, 100, 100)])
    
    # Fitur 2: Jumlah Bytes yang dikirim
    # Normal: kecil/sedang, Serangan (Exfiltration): besar sekali
    src_bytes = np.concatenate([np.random.randint(100, 2000, 900), np.random.randint(10000, 50000, 100)])
    
    # Fitur 3: Protocol (0=TCP, 1=UDP, 2=ICMP)
    protocol = np.random.randint(0, 3, n_samples)
    
    # Label: 0 = Normal, 1 = Attack
    # Kita buat logika sederhana: Jika durasi lama ATAU bytes besar -> kemungkinan Attack
    labels = []
    for d, b in zip(duration, src_bytes):
        if d > 10 or b > 8000:
            labels.append("Attack")
        else:
            labels.append("Normal")
            
    df = pd.DataFrame({
        'Duration_Sec': duration,
        'Src_Bytes': src_bytes,
        'Protocol_Code': protocol,
        'Label': labels
    })
    
    return df

# --- BAGIAN 2: TRAINING MODEL AI ---
print("Sedang membuat dataset & melatih AI...")
df = create_dummy_dataset()

# Pisahkan fitur (X) dan label (y)
X = df[['Duration_Sec', 'Src_Bytes', 'Protocol_Code']]
y = df['Label']

# Train Model Random Forest sederhana
clf = RandomForestClassifier(n_estimators=10, random_state=42)
clf.fit(X, y)
print("Training selesai!")

# --- BAGIAN 3: FUNGSI PREDIKSI UNTUK DEMO ---
def predict_traffic(duration, bytes_val, protocol_name):
    # Konversi nama protokol ke kode angka
    proto_map = {"TCP": 0, "UDP": 1, "ICMP": 2}
    proto_code = proto_map.get(protocol_name, 0)
    
    # Prediksi
    input_data = [[duration, bytes_val, proto_code]]
    prediction = clf.predict(input_data)[0]
    probs = clf.predict_proba(input_data)[0] # Ambil probabilitas
    
    # Format Output
    prob_normal = probs[0] if clf.classes_[0] == 'Normal' else probs[1]
    prob_attack = probs[1] if clf.classes_[0] == 'Normal' else probs[0]
    
    # Siapkan Chart Data
    chart_data = {"Normal": prob_normal, "Attack": prob_attack}
    
    if prediction == "Attack":
        msg = f"⚠️ TERDETEKSI ANOMALI!\nKoneksi ini mencurigakan (Durasi/Bytes tidak wajar)."
    else:
        msg = f"✅ LALU LINTAS NORMAL\nKoneksi terlihat wajar."
        
    return msg, chart_data, df.head(10) # Tampilkan juga contoh data tabel

# --- BAGIAN 4: UI GRADIO ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 📊 AI Network Traffic Analyzer (Modul 2)")
    gr.Markdown("Demo ini menggunakan **Random Forest** yang dilatih pada dataset simulasi lalu lintas jaringan.")
    
    with gr.Row():
        with gr.Column():
            gr.Markdown("### 1. Atur Parameter Paket Data")
            
            # Input Slider untuk simulasi serangan
            in_duration = gr.Slider(0, 100, value=2, label="Durasi Koneksi (detik)", 
                                    info="Koneksi normal biasanya singkat. Serangan DoS biasanya lama.")
            
            in_bytes = gr.Slider(0, 50000, value=500, label="Ukuran Data (Bytes)", 
                                 info="Ukuran paket wajar vs Indikasi pencurian data (Exfiltration).")
            
            in_proto = gr.Dropdown(["TCP", "UDP", "ICMP"], value="TCP", label="Protokol")
            
            btn = gr.Button("🔍 Analisis Trafik", variant="primary")

        with gr.Column():
            gr.Markdown("### 2. Hasil Analisis AI")
            out_text = gr.Textbox(label="Kesimpulan")
            out_label = gr.Label(label="Probabilitas")
    
    # Tambahan: Menampilkan Dataset agar siswa paham bentuk datanya
    gr.Markdown("### 3. Intip Isi Dataset (10 Baris Pertama)")
    gr.Markdown("Ini adalah contoh data mentah yang dipelajari oleh AI:")
    out_df = gr.Dataframe(label="Sample Dataset Log")

    btn.click(fn=predict_traffic, 
              inputs=[in_duration, in_bytes, in_proto], 
              outputs=[out_text, out_label, out_df])

# Launch
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)