api/main.py

"""
ImageVerify AI — Production Inference API
=========================================
FastAPI server for the MFFT model with:
  - Tiered rate limiting (free/pro/enterprise)
  - Explainable predictions with heatmap visualization
  - Batch processing
  - Report generation (JSON/PDF)
"""

import io
import json
import time
import base64
from pathlib import Path
from datetime import datetime
from typing import List, Optional
from contextlib import asynccontextmanager

import torch
import torch.nn.functional as F
import numpy as np
from PIL import Image
from fastapi import FastAPI, UploadFile, File, HTTPException, Depends, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, Response
from pydantic import BaseModel, Field

from .schemas import (
    PredictionRequest, PredictionResponse, BatchPredictionResponse,
    SinglePrediction, HealthResponse, UsageStats, ErrorResponse
)
from .model_server import ModelServer


import os

# variant registry: every loaded variant is selectable per request
MODEL_INFO = {
    "tiny":  {"params": "372K",  "description": "Fastest - edge/mobile profile"},
    "base":  {"params": "1.62M", "description": "Balanced accuracy and speed"},
    "large": {"params": "6.30M", "description": "Highest accuracy profile"},
}
model_registry: dict = {}
DEFAULT_VARIANT = os.environ.get("MFFT_VARIANT", "base")
IMAGE_SIZE = int(os.environ.get("MFFT_IMAGE_SIZE", "384"))
DEMO_MODE = os.environ.get("MFFT_DEMO") == "1"


def _find_checkpoint(variant: str) -> Optional[str]:
    checkpoint_dir = Path(__file__).parent.parent / "model" / "checkpoints"
    env_dir = os.environ.get("MFFT_CHECKPOINT_DIR")
    candidates = []
    if env_dir:
        candidates += [
            str(Path(env_dir) / f"{variant}.pt"),
            str(Path(env_dir) / f"best_mfft_{variant}.pt"),
        ]
    candidates += [
        # written by the train notebooks (full-scale run)
        str(checkpoint_dir / f"{variant}_model" / f"best_mfft_{variant}.pt"),
        # written by the verification / pilot notebooks
        str(checkpoint_dir / "verify" / f"{variant}_model" / "best.pt"),
        str(checkpoint_dir / "test" / f"{variant}_model" / "best.pt"),
        # legacy locations
        str(checkpoint_dir / f"best_mfft_{variant}.pt"),
    ]
    return next((c for c in candidates if Path(c).exists()), None)


@asynccontextmanager
async def lifespan(app: FastAPI):
    env_ckpt = os.environ.get("MFFT_CHECKPOINT")  # single-model override
    for variant in MODEL_INFO:
        ckpt = env_ckpt if (env_ckpt and variant == DEFAULT_VARIANT) \
            else _find_checkpoint(variant)
        if ckpt:
            try:
                model_registry[variant] = ModelServer(
                    ckpt, variant=variant, image_size=IMAGE_SIZE)
            except Exception as e:
                print(f"[startup] {variant}: failed to load {ckpt}: {e}")
        else:
            print(f"[startup] {variant}: no checkpoint found, not serving")

    if not model_registry:
        if os.environ.get("MFFT_ALLOW_RANDOM") == "1":
            model_registry[DEFAULT_VARIANT] = ModelServer(
                None, variant=DEFAULT_VARIANT, image_size=IMAGE_SIZE)
        else:
            raise RuntimeError(
                "No MFFT checkpoint found for any variant - refusing to serve "
                "random predictions. Train the model first (see DGX_RUN_GUIDE.md), "
                "set MFFT_CHECKPOINT_DIR, or set MFFT_ALLOW_RANDOM=1 (dev only)."
            )
    print(f"[startup] serving variants: {sorted(model_registry)}")
    yield
    model_registry.clear()


def get_server(model: str) -> ModelServer:
    variant = (model or DEFAULT_VARIANT).lower()
    if variant not in model_registry:
        raise HTTPException(
            404,
            f"Model '{variant}' not available. Loaded: {sorted(model_registry)}",
        )
    return model_registry[variant]


app = FastAPI(
    title="ImageVerify AI",
    description="AI-Generated Image Detection API",
    version="2.0.0",
    lifespan=lifespan,
)

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


TIER_LIMITS = {
    "free": {"rpm": 10, "batch_size": 1, "report": False},
    "pro": {"rpm": 100, "batch_size": 10, "report": True},
    "enterprise": {"rpm": 1000, "batch_size": 100, "report": True},
}


class UsageTracker:
    def __init__(self):
        self.requests: dict = {}

    def check_limit(self, api_key: str) -> bool:
        tier = self._get_tier(api_key)
        limit = TIER_LIMITS[tier]
        now = time.time()
        minute_ago = now - 60

        if api_key not in self.requests:
            self.requests[api_key] = []

        self.requests[api_key] = [
            t for t in self.requests[api_key] if t > minute_ago
        ]

        if len(self.requests[api_key]) >= limit["rpm"]:
            raise HTTPException(
                status_code=429,
                detail=f"Rate limit exceeded ({limit['rpm']} req/min for {tier} tier)",
            )

        self.requests[api_key].append(now)
        return True

    def _get_tier(self, api_key: str) -> str:
        if not api_key or api_key == "free":
            return "free"
        if api_key.startswith("pro_"):
            return "pro"
        if api_key.startswith("ent_"):
            return "enterprise"
        return "free"

    def get_tier_limits(self, api_key: str) -> dict:
        return TIER_LIMITS[self._get_tier(api_key)]


usage_tracker = UsageTracker()


def get_api_key(authorization: str = "") -> str:
    if authorization.startswith("Bearer "):
        return authorization[7:]
    return "free"


@app.get("/", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="healthy",
        model_loaded=len(model_registry) > 0,
        version="2.1.0",
        timestamp=datetime.now().isoformat(),
    )


@app.get("/models")
async def list_models():
    """Available model variants for the `model` query parameter."""
    return {
        "default": DEFAULT_VARIANT if DEFAULT_VARIANT in model_registry
        else (sorted(model_registry)[0] if model_registry else None),
        "models": [
            {
                "id": v,
                "loaded": v in model_registry,
                **MODEL_INFO[v],
            }
            for v in MODEL_INFO
        ],
    }


@app.post("/predict", response_model=PredictionResponse)
async def predict(
    file: UploadFile = File(...),
    model: str = "base",
    api_key: str = Depends(get_api_key),
):
    usage_tracker.check_limit(api_key)
    tier_limits = usage_tracker.get_tier_limits(api_key)
    server = get_server(model)

    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(400, "File must be an image")

    contents = await file.read()
    if len(contents) > 20 * 1024 * 1024:
        raise HTTPException(400, "File too large (max 20MB)")

    try:
        image = Image.open(io.BytesIO(contents)).convert("RGB")
    except Exception:
        raise HTTPException(400, "Invalid image file")

    result = server.predict(image)

    response = PredictionResponse(
        prediction="ai_generated" if result["prediction"] == 1 else "real",
        confidence=round(float(result["confidence"]), 4),
        real_probability=round(float(result["real_prob"]), 4),
        ai_probability=round(float(result["ai_prob"]), 4),
        processing_time_ms=round(result["processing_time_ms"], 2),
        tier=tier_limits,
    )

    if tier_limits["report"] or DEMO_MODE:
        heatmap_b64 = _heatmap_to_base64(result["heatmaps"])
        response.anomaly_heatmap = heatmap_b64
        response.frequency_band_contributions = result.get("frequency_band_contributions", {})

    return response


@app.post("/predict/batch", response_model=BatchPredictionResponse)
async def predict_batch(
    files: List[UploadFile] = File(...),
    model: str = "base",
    api_key: str = Depends(get_api_key),
):
    usage_tracker.check_limit(api_key)
    tier_limits = usage_tracker.get_tier_limits(api_key)
    server = get_server(model)

    if len(files) > tier_limits["batch_size"]:
        raise HTTPException(
            400,
            f"Batch limit exceeded (max {tier_limits['batch_size']} for your tier)",
        )

    results = []
    for file in files:
        contents = await file.read()
        try:
            image = Image.open(io.BytesIO(contents)).convert("RGB")
            result = server.predict(image)
            results.append(SinglePrediction(
                filename=file.filename or "unknown",
                prediction="ai_generated" if result["prediction"] == 1 else "real",
                confidence=round(float(result["confidence"]), 4),
                real_probability=round(float(result["real_prob"]), 4),
                ai_probability=round(float(result["ai_prob"]), 4),
                processing_time_ms=round(result["processing_time_ms"], 2),
            ))
        except Exception as e:
            results.append(SinglePrediction(
                filename=file.filename or "unknown",
                prediction="error",
                confidence=0.0,
                processing_time_ms=0,
                error=str(e),
            ))

    avg_real = np.mean([r.real_probability for r in results if r.real_probability])
    avg_ai = np.mean([r.ai_probability for r in results if r.ai_probability])
    ai_count = sum(1 for r in results if r.prediction == "ai_generated")
    real_count = sum(1 for r in results if r.prediction == "real")

    return BatchPredictionResponse(
        results=results,
        summary={
            "total": len(results),
            "ai_generated": ai_count,
            "real": real_count,
            "avg_real_probability": round(float(avg_real), 4),
            "avg_ai_probability": round(float(avg_ai), 4),
        },
        tier=tier_limits,
    )


@app.get("/usage", response_model=UsageStats)
async def get_usage(api_key: str = Depends(get_api_key)):
    tier = usage_tracker._get_tier(api_key)
    limits = TIER_LIMITS[tier]
    return UsageStats(
        tier=tier,
        requests_this_minute=len(usage_tracker.requests.get(api_key, [])),
        rate_limit=limits["rpm"],
    )


def _heatmap_to_base64(heatmaps: torch.Tensor) -> str:
    if heatmaps is None:
        return ""
    h = heatmaps[0].cpu().numpy()
    h = (h - h.min()) / (h.max() - h.min() + 1e-8)
    h = (h * 255).astype(np.uint8)
    img = Image.fromarray(h[0]) if h.ndim == 3 else Image.fromarray(h)
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode()



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