app.py

"""
app.py
Main FastAPI application — background removal SaaS backend.
Updated for Hugging Face Spaces with CORS, API key auth, and rate limiting.
"""

import logging
import os
import time
from collections import defaultdict
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, Request, HTTPException, Depends, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse, JSONResponse
from fastapi.exceptions import RequestValidationError
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials

from config.constants import ENGINE_TAG, API_VERSION
from config.settings import ALLOWED_ORIGINS
from models.schemas import HealthResponse, RootResponse
from services.image_processor import (
    load_model,
    warm_model,
    is_model_ready,
    remove_background,
    image_to_png_buffer,
)
from services.validators import run_all_validations

# ── Logging ────────────────────────────────────────────────────────────────────
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-8s | %(name)s | %(message)s",
)
logger = logging.getLogger(__name__)

# ── Security & Rate Limiting ───────────────────────────────────────────────────
security = HTTPBearer(auto_error=False)  # API key is optional

# Get API key from Hugging Face Secrets (set in Space Settings)
API_KEY = os.environ.get("API_SECRET_KEY")
# Rate limiting storage (in-memory, resets on container restart)
rate_limits = defaultdict(list)
RATE_LIMIT_REQUESTS = int(os.environ.get("RATE_LIMIT_REQUESTS", 100))   # 100 requests
RATE_LIMIT_PERIOD = int(os.environ.get("RATE_LIMIT_PERIOD", 86400))      # per day (seconds)

def verify_api_key(credentials: HTTPAuthorizationCredentials = Depends(security)):
    """
    Optional API key verification.
    If API_KEY is set in environment, requests must provide it.
    If not set, all requests are allowed (useful for testing).
    """
    if API_KEY is None:
        # No key configured → allow all requests (open API)
        return True
    if credentials is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="API key required. Please provide Authorization: Bearer <key>",
            headers={"WWW-Authenticate": "Bearer"},
        )
    if credentials.credentials != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail="Invalid API key",
        )
    return True

def rate_limit(request: Request):
    """Simple rate limiting by IP address."""
    if API_KEY is None:
        # No key -> no rate limiting (or you can still apply it)
        return True
    client_ip = request.client.host
    now = time.time()
    # Clean old entries
    rate_limits[client_ip] = [ts for ts in rate_limits[client_ip] if now - ts < RATE_LIMIT_PERIOD]
    if len(rate_limits[client_ip]) >= RATE_LIMIT_REQUESTS:
        raise HTTPException(
            status_code=status.HTTP_429_TOO_MANY_REQUESTS,
            detail=f"Rate limit exceeded: {RATE_LIMIT_REQUESTS} requests per {RATE_LIMIT_PERIOD//3600} hours",
        )
    rate_limits[client_ip].append(now)
    return True

# ── Startup / Shutdown lifecycle ───────────────────────────────────────────────
@asynccontextmanager
async def lifespan(app: FastAPI):
    # --- startup ---
    logger.info("=== Background Remover API starting ===")
    logger.info(f"API Key enabled: {API_KEY is not None}")
    logger.info(f"Rate limit: {RATE_LIMIT_REQUESTS} requests per {RATE_LIMIT_PERIOD//3600} hours per IP")
    load_model()
    warm_model()
    logger.info("=== API ready — accepting requests ===")
    yield
    # --- shutdown ---
    logger.info("=== Background Remover API shutting down ===")

# ── Application factory ────────────────────────────────────────────────────────
app = FastAPI(
    title="Background Remover API",
    version=API_VERSION,
    description="Remove image backgrounds in under a second using the Silueta ONNX model.",
    lifespan=lifespan,
)

# ── CORS ───────────────────────────────────────────────────────────────────────
# ALLOWED_ORIGINS is read from environment variable (comma-separated)
# Example: https://your-frontend.vercel.app,http://localhost:5173
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["GET", "POST"],
    allow_headers=["*"],
)

# ── Exception Handlers ─────────────────────────────────────────────────────────
@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    return JSONResponse(status_code=422, content={"detail": "Invalid request payload."})

@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    logger.exception("Unhandled exception")
    return JSONResponse(status_code=500, content={"detail": "Internal server error."})

# ── Routes ─────────────────────────────────────────────────────────────────────
@app.get("/", response_model=RootResponse, tags=["Meta"])
async def root():
    return RootResponse()

@app.get("/health", response_model=HealthResponse, tags=["Meta"])
async def health():
    from config.constants import MODEL_UNAVAILABLE
    if not is_model_ready():
        raise HTTPException(status_code=503, detail=MODEL_UNAVAILABLE)
    return HealthResponse(
        status="healthy",
        engine=ENGINE_TAG,
        model_loaded=True,
    )

@app.post(
    "/api/v1/remove-background",
    tags=["Processing"],
    dependencies=[Depends(verify_api_key), Depends(rate_limit)],  # <-- Added auth & rate limit
    responses={
        200: {"content": {"image/png": {}}, "description": "Transparent PNG stream"},
        400: {"description": "Corrupted image / bad dimensions"},
        401: {"description": "Missing or invalid API key"},
        413: {"description": "File too large"},
        415: {"description": "Unsupported file type"},
        429: {"description": "Rate limit exceeded"},
        500: {"description": "Model processing failure"},
        503: {"description": "Inference engine unavailable"},
    },
)
async def remove_background_endpoint(
    file: UploadFile = File(..., description="JPEG, PNG, or WEBP image ≤ 5 MB"),
):
    logger.info("Request received: %s (%s)", file.filename, file.content_type)

    # Validation
    _raw_bytes, image = await run_all_validations(file)
    logger.info("Validation passed for %s", file.filename)

    # Inference
    logger.info("Processing started.")
    result_image = remove_background(image)
    logger.info("Processing complete.")

    # Export
    png_buffer = image_to_png_buffer(result_image)

    return StreamingResponse(
        png_buffer,
        media_type="image/png",
        headers={
            "Content-Disposition": 'attachment; filename="result.png"',
            "X-Engine": ENGINE_TAG,
        },
    )