from fastapi import FastAPI, Request
from pydantic import BaseModel
import torch
from src.model import AnomalyDetector, detect_anomaly
import logging

app = FastAPI(title="NetGuard-AI API", description="Real-Time Network Intrusion Detection API")
logger = logging.getLogger("netguard-api")

# Load Model
model = AnomalyDetector()
try:
    model.load_state_dict(torch.load('models/autoencoder.pth', map_location='cpu'))
    print("Model loaded successfully.")
except FileNotFoundError:
    print("Warning: Model file not found. Using untrained model.")
model.eval()

class TrafficData(BaseModel):
    features: list[float] # Expected length 41

@app.post("/predict")
async def predict(data: TrafficData):
    """
    Receives a single network flow and predicts if it's anomalous.
    """
    if len(data.features) != 41:
        return {"error": "Invalid feature length. Expected 41."}
    
    tensor_data = torch.tensor([data.features], dtype=torch.float32)
    is_anomaly, score = detect_anomaly(model, tensor_data, threshold=0.5)
    
    result = {
        "is_anomaly": bool(is_anomaly.item()),
        "anomaly_score": float(score.item()),
        "status": "Blocked" if is_anomaly.item() else "Allowed"
    }
    
    if result["is_anomaly"]:
        logger.warning(f"Intrusion Detected! Score: {result['anomaly_score']}")
        
    return result

@app.get("/health")
def health_check():
    return {"status": "healthy"}