app.py

"""
app.py — FastAPI application for the Gridlock Traffic Violation API.
=====================================================================
Endpoints:
  POST /detect         Upload image → run pipeline → return violations + annotated image
  GET  /health         Health check (models loaded?)
  GET  /stats          Aggregate violation statistics
  GET  /violations     Search/filter past violations
"""

import io
import base64
import logging
import tempfile
from pathlib import Path
from contextlib import asynccontextmanager

import cv2
import numpy as np
from fastapi import FastAPI, File, UploadFile, Query, HTTPException
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware

from config import API_TITLE, API_VERSION, OUTPUT_DIR
from pipeline import ParallelDetectionPipeline
from annotations import annotate_from_pipeline_result
from storage import init_db, store_result, get_statistics, search_violations

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(name)s] %(levelname)s: %(message)s",
)
logger = logging.getLogger("gridlock.app")

# ── Global pipeline instance (loaded once at startup) ────────────────────────
_pipeline: ParallelDetectionPipeline = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Load models on startup, cleanup on shutdown."""
    global _pipeline
    logger.info("Starting Gridlock API — loading models...")
    init_db()
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    _pipeline = ParallelDetectionPipeline()
    logger.info("Models loaded. API ready.")
    yield
    logger.info("Shutting down Gridlock API.")
    _pipeline = None


app = FastAPI(
    title=API_TITLE,
    version=API_VERSION,
    description=(
        "AI-powered traffic violation detection API. "
        "Upload an image to detect helmet violations, seatbelt violations, "
        "wrong-way driving, illegal parking, triple riding, and more."
    ),
    lifespan=lifespan,
)

# CORS for frontend integration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ═══════════════════════════════════════════════════════════════════════════════
# ENDPOINTS
# ═══════════════════════════════════════════════════════════════════════════════

@app.get("/health")
async def health_check():
    """Check if the API is running and models are loaded."""
    return {
        "status": "healthy" if _pipeline is not None else "loading",
        "models_loaded": _pipeline is not None,
    }


@app.post("/detect")
async def detect_violations(
    file: UploadFile = File(...),
    return_annotated_image: bool = Query(True, description="Include base64 annotated image in response"),
    save_annotated: bool = Query(True, description="Save annotated image to disk"),
):
    """
    Upload an image and detect traffic violations.

    Returns violations list, vehicle counts, processing time,
    and optionally the annotated image as base64.
    """
    if _pipeline is None:
        raise HTTPException(status_code=503, detail="Models are still loading. Try again shortly.")

    # Validate file type
    if not file.filename:
        raise HTTPException(status_code=400, detail="No file provided.")

    suffix = Path(file.filename).suffix.lower()
    if suffix not in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type '{suffix}'. Use: jpg, jpeg, png, bmp, webp.",
        )

    # Read uploaded image into a temp file (Roboflow SDK needs a file path)
    contents = await file.read()
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(
            suffix=suffix, delete=False, dir=str(OUTPUT_DIR)
        ) as tmp:
            tmp.write(contents)
            tmp_path = tmp.name

        # Run the pipeline
        result = _pipeline.process(tmp_path, return_annotations=True)

        # Generate annotated image
        img = cv2.imread(tmp_path)
        annotated_img = None
        annotated_path = None

        if img is not None and (return_annotated_image or save_annotated):
            annotated_img = annotate_from_pipeline_result(img, result)

            if save_annotated:
                out_name = Path(file.filename).stem + "_annotated.jpg"
                annotated_path = str(OUTPUT_DIR / out_name)
                cv2.imwrite(annotated_path, annotated_img, [cv2.IMWRITE_JPEG_QUALITY, 95])

        # Store in database
        store_result(result, file.filename, annotated_path)

        # Build response
        response = {
            "image_id": result["image_id"],
            "timestamp": result["timestamp"],
            "processing_time_ms": result["processing_time_ms"],
            "vehicles_detected": result["vehicles_detected"],
            "violations": result["violations"],
            "illegal_parking": result.get("illegal_parking", []),
        }

        if return_annotated_image and annotated_img is not None:
            _, buffer = cv2.imencode(".jpg", annotated_img, [cv2.IMWRITE_JPEG_QUALITY, 90])
            response["annotated_image_base64"] = base64.b64encode(buffer).decode("utf-8")

        if annotated_path:
            response["annotated_image_path"] = annotated_path

        # Remove annotation_data from response (internal use only)
        response.pop("annotation_data", None)

        return JSONResponse(content=response)

    except Exception as e:
        logger.exception("Error processing image: %s", e)
        raise HTTPException(status_code=500, detail=f"Processing failed: {str(e)}")

    finally:
        # Clean up temp file
        if tmp_path:
            try:
                Path(tmp_path).unlink(missing_ok=True)
            except Exception:
                pass


@app.get("/stats")
async def violation_stats():
    """Get aggregate violation statistics."""
    try:
        stats = get_statistics()
        return JSONResponse(content=stats)
    except Exception as e:
        logger.exception("Error fetching stats: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/violations")
async def list_violations(
    plate: str = Query(None, description="Filter by license plate (partial match)"),
    violation_type: str = Query(None, description="Filter by violation type (e.g., 'no_helmet', 'wrong_side')"),
    start_date: str = Query(None, description="Start date (ISO format)"),
    end_date: str = Query(None, description="End date (ISO format)"),
    limit: int = Query(50, ge=1, le=200, description="Max results"),
    offset: int = Query(0, ge=0, description="Offset for pagination"),
):
    """Search and filter violation records."""
    try:
        results = search_violations(
            plate=plate,
            violation_type=violation_type,
            start_date=start_date,
            end_date=end_date,
            limit=limit,
            offset=offset,
        )
        return JSONResponse(content={"count": len(results), "violations": results})
    except Exception as e:
        logger.exception("Error searching violations: %s", e)
        raise HTTPException(status_code=500, detail=str(e))


# ═══════════════════════════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=False)