backend/main.py

"""FastAPI backend for DeepShield detection API.

Endpoints:
  POST /analyze        — upload image file, get EnsembleResult JSON
  GET  /health         — liveness probe
  GET  /detectors      — list active detectors and their config
  GET  /history        — get analysis history
  GET  /history/{id}   — get analysis by ID
  GET  /history/stats  — get summary statistics
  DELETE /history/{id} — delete analysis by ID

Run:
  uvicorn backend.main:app --reload --port 8002
"""

from __future__ import annotations

import asyncio
import logging
import tempfile

logging.basicConfig(level=logging.INFO, format="%(levelname)s [%(name)s] %(message)s")
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import json

from fastapi import FastAPI, File, HTTPException, UploadFile, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.encoders import jsonable_encoder
from fastapi.responses import StreamingResponse

from backend.core.pipeline import DetectionPipeline
from backend.core.schema import EnsembleResult
from backend.db.mongodb import connect_db, close_db, get_db
from backend.crud.analysis import (
    save_analysis,
    get_analysis_by_id,
    get_analysis_history,
    get_history_stats,
    delete_analysis,
)

app = FastAPI(
    title="DeepShield Detection API",
    description="Multi-API ensemble AI image detector with history",
    version="2.0.0",
)

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

_pipelines: dict[str, DetectionPipeline] = {}

_ALLOWED_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"}


def _warmup_hybrid_model() -> None:
    """Pre-load PyTorch hybrid model weights in a background thread at startup."""
    from backend.core.config import settings
    if not settings.enable_hybrid_model:
        return
    try:
        from backend.detectors.hybrid_model import _load_bundle, _model_bundles, _load_errors
        for version in ("latest",):
            if version not in _model_bundles:
                bundle, err = _load_bundle(version)
                _model_bundles[version] = bundle
                _load_errors[version] = err
                if err:
                    print(f"  [hybrid_model] warmup skipped ({version}): {err}")
                else:
                    print(f"  [hybrid_model] warmed up ({version})")
    except Exception as e:
        print(f"  [hybrid_model] warmup error: {e}")


@app.on_event("startup")
async def startup():
    """Initialize pipeline cache, warm up hybrid model, and connect MongoDB."""
    for version in ("latest", "backup"):
        _pipelines[version] = DetectionPipeline(model_version=version)

    # Warm up hybrid model in background thread (non-blocking)
    loop = asyncio.get_event_loop()
    loop.run_in_executor(ThreadPoolExecutor(max_workers=1), _warmup_hybrid_model)

    try:
        await connect_db()
    except Exception as e:
        print(f"WARNING: MongoDB connection failed: {e}")
        print("  Running without history persistence")


@app.on_event("shutdown")
async def shutdown():
    """Close MongoDB on app shutdown."""
    await close_db()


def _get_pipeline(model_version: str = "latest") -> DetectionPipeline:
    return _pipelines.get(model_version) or _pipelines["latest"]


@app.get("/health")
async def health():
    return {"status": "ok", "version": "2.0.0"}


@app.get("/detectors")
async def list_detectors():
    pipeline = _get_pipeline()
    return {"active_detectors": pipeline.active_detectors}


@app.get("/models")
async def list_models():
    """List available hybrid model versions."""
    return {
        "available_models": ["latest", "backup"],
        "default_model": "latest",
        "description": "latest: most recent trained model, backup: previous stable version"
    }


# Heartbeat cadence for the streamed /analyze response. Analysis runs 30s+;
# mobile carrier NAT drops connections idle for ~30s, so we emit a whitespace
# byte on this interval to keep the connection active. Leading whitespace is
# ignored by JSON parsers, so the client can parse the final payload directly.
_ANALYZE_HEARTBEAT_SECONDS = 5.0


@app.post("/analyze")
async def analyze(
    file: UploadFile = File(...),
    model_version: str = Query("latest", pattern="^(latest|backup)$"),
    detectors: str | None = Query(None, description="Comma-separated detector names to use. Omit for all enabled detectors."),
    fusion_strategy: str | None = Query(None, pattern="^(weighted|voting)$", description="Fusion strategy: 'weighted' (Bayesian log-odds) or 'voting' (majority vote). Overrides server default.")
) -> StreamingResponse:
    ext = Path(file.filename or "").suffix.lower()
    if ext not in _ALLOWED_EXTENSIONS:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported file type '{ext}'. Allowed: {_ALLOWED_EXTENSIONS}",
        )

    enabled_detectors = (
        [d.strip() for d in detectors.split(",") if d.strip()]
        if detectors else None
    )

    with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as tmp:
        content = await file.read()
        tmp.write(content)
        tmp_path = tmp.name

    filename = file.filename or "unknown"

    async def _run() -> EnsembleResult:
        pipeline = _get_pipeline(model_version=model_version)
        result = await pipeline.run(tmp_path, enabled_detectors=enabled_detectors, fusion_strategy=fusion_strategy)
        try:
            db = get_db()
            await save_analysis(db, filename, result)
        except Exception as db_err:
            logging.getLogger(__name__).warning("MongoDB save failed: %s", db_err)
        return result

    async def stream():
        # The connection is committed (200 + headers) once streaming starts, so
        # errors after this point are delivered as a JSON {"error": ...} payload
        # rather than an HTTP error status.
        task = asyncio.create_task(_run())
        try:
            while True:
                done, _ = await asyncio.wait({task}, timeout=_ANALYZE_HEARTBEAT_SECONDS)
                if done:
                    break
                yield b" "  # heartbeat — keeps mobile carrier NAT from dropping the connection
            result = task.result()
            yield json.dumps(jsonable_encoder(result)).encode()
        except FileNotFoundError as e:
            yield json.dumps({"error": str(e)}).encode()
        except Exception as e:
            logging.getLogger(__name__).exception("Pipeline error")
            yield json.dumps({"error": str(e)}).encode()
        finally:
            Path(tmp_path).unlink(missing_ok=True)

    return StreamingResponse(stream(), media_type="application/json")


@app.get("/history")
async def get_history(limit: int = Query(100, ge=1, le=1000), skip: int = Query(0, ge=0)):
    """Get analysis history with pagination."""
    try:
        db = get_db()
        history = await get_analysis_history(db, limit=limit, skip=skip)
        # Convert ObjectId to string for JSON
        for item in history:
            item["_id"] = str(item["_id"])
        return {"results": history, "limit": limit, "skip": skip}
    except RuntimeError:
        raise HTTPException(status_code=503, detail="Database not available")


@app.get("/history/stats")
async def get_stats():
    """Get summary statistics of all analyses."""
    try:
        db = get_db()
        stats = await get_history_stats(db)
        return {"stats": stats}
    except RuntimeError:
        raise HTTPException(status_code=503, detail="Database not available")


@app.get("/history/{analysis_id}")
async def get_analysis(analysis_id: str):
    """Get specific analysis by ID."""
    try:
        db = get_db()
        result = await get_analysis_by_id(db, analysis_id)
        if not result:
            raise HTTPException(status_code=404, detail="Analysis not found")
        result["_id"] = str(result["_id"])
        return result
    except RuntimeError:
        raise HTTPException(status_code=503, detail="Database not available")


@app.delete("/history/{analysis_id}")
async def delete_analysis_endpoint(analysis_id: str):
    """Delete analysis by ID."""
    try:
        db = get_db()
        deleted = await delete_analysis(db, analysis_id)
        if not deleted:
            raise HTTPException(status_code=404, detail="Analysis not found")
        return {"message": "Analysis deleted"}
    except RuntimeError:
        raise HTTPException(status_code=503, detail="Database not available")