from fastapi import FastAPI
from pydantic import BaseModel
import joblib
import pandas as pd

app = FastAPI()

# Load your model
model = joblib.load('cicids_model.pkl')

# Features your model expects
FEATURES = [
    ' Flow Duration', ' Flow Packets/s', 'Flow Bytes/s',
    ' Total Fwd Packets', ' Total Backward Packets',
    ' Packet Length Mean', ' Packet Length Std', ' Flow IAT Mean',
    ' Flow IAT Std', ' SYN Flag Count', ' ACK Flag Count',
    ' RST Flag Count', ' Average Packet Size', ' Down/Up Ratio'
]

class NetworkFlowData(BaseModel):
    flow_duration: float
    flow_packets_per_sec: float
    flow_bytes_per_sec: float
    total_fwd_packets: float
    total_backward_packets: float
    packet_length_mean: float
    packet_length_std: float
    flow_iat_mean: float
    flow_iat_std: float
    syn_flag_count: float
    ack_flag_count: float
    rst_flag_count: float
    average_packet_size: float
    down_up_ratio: float

@app.get("/")
def root():
    return {
        "message": "CICIDS2017 Network Attack Detection API",
        "attack_types": list(model.classes_),
        "features_expected": FEATURES
    }

@app.post("/predict")
def predict(data: NetworkFlowData):
    input_dict = {
        ' Flow Duration': data.flow_duration,
        ' Flow Packets/s': data.flow_packets_per_sec,
        'Flow Bytes/s': data.flow_bytes_per_sec,
        ' Total Fwd Packets': data.total_fwd_packets,
        ' Total Backward Packets': data.total_backward_packets,
        ' Packet Length Mean': data.packet_length_mean,
        ' Packet Length Std': data.packet_length_std,
        ' Flow IAT Mean': data.flow_iat_mean,
        ' Flow IAT Std': data.flow_iat_std,
        ' SYN Flag Count': data.syn_flag_count,
        ' ACK Flag Count': data.ack_flag_count,
        ' RST Flag Count': data.rst_flag_count,
        ' Average Packet Size': data.average_packet_size,
        ' Down/Up Ratio': data.down_up_ratio
    }
    input_df = pd.DataFrame([input_dict])
    prediction = model.predict(input_df)[0]
    probabilities = model.predict_proba(input_df)[0]
    confidence = float(max(probabilities))
    
    return {
        "prediction": str(prediction),
        "confidence": confidence
    }