app/main.py

from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.cors import CORSMiddleware
import shutil
import os
from app.utils import predict_image, get_ip_location, fetch_hospitals, ALERT_CLASSES

app = FastAPI(
    title="AccidentAI API",
    description="YOLOv8-powered accident & fire detection with emergency hospital alerts",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

UPLOAD_FOLDER = "temp"
os.makedirs(UPLOAD_FOLDER, exist_ok=True)


@app.get("/")
def home():
    return {"message": "AccidentAI — YOLO Accident Detection API 🚀", "status": "running"}


@app.post("/predict")
async def predict(file: UploadFile = File(...), lat: float = None, lon: float = None):
    """
    Upload an image for accident/fire detection.
    Optionally pass lat/lon for hospital lookup; otherwise IP geolocation is used.
    """
    file_path = os.path.join(UPLOAD_FOLDER, file.filename)

    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    detections = predict_image(file_path)

    # Clean up uploaded file
    try:
        os.remove(file_path)
    except OSError:
        pass

    # Check for alerts
    alert_labels = [d["label"] for d in detections if d["label"] in ALERT_CLASSES]
    has_alert = len(alert_labels) > 0

    response = {
        "filename": file.filename,
        "detections": detections,
        "alert": has_alert,
        "alert_labels": list(set(alert_labels)),
    }

    # If accident/fire detected → find nearby hospitals
    if has_alert:
        if lat is not None and lon is not None:
            location = {"lat": lat, "lon": lon}
        else:
            location = get_ip_location()

        if location:
            hospitals = fetch_hospitals(location["lat"], location["lon"])
            response["location"] = location
            response["hospitals_nearby"] = hospitals
            response["hospitals_notified"] = len(hospitals[:3])
            response["nearest_hospital"] = hospitals[0] if hospitals else None
        else:
            response["location"] = None
            response["hospitals_nearby"] = []

    return response