"""
Patch-MFFT Tampering Forensics API
==================================
FastAPI service providing:
- 3-Class Forensics: Authentic vs Full-AI vs Spliced/Inpainted Tampering
- Real-time Pixel-Level Tampering Localization Masks (Base64 PNG)
- Radial Fourier Spectral Band Energy Decomposition
"""

import io
import os
import hashlib
from collections import OrderedDict
from datetime import datetime, timezone
from pathlib import Path
from typing import Dict, List, Optional
from contextlib import asynccontextmanager

from PIL import Image
from fastapi import FastAPI, UploadFile, File, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

from .schemas import HealthResponse, ModelsResponse, PredictionResponse, ModelInfo
from .model_server import TamperingModelServer, predict_ensemble

MODEL_INFO = {
    "tiny":     {"params": "563K",   "description": "Fastest - edge/mobile tampering detection profile (~4 ms)"},
    "base":     {"params": "3.28M",  "description": "Standard Profile - balanced accuracy and localized inpainting"},
    "large":    {"params": "13.94M", "description": "Deep Profile - recommended for vintage/retro AI & degraded media"},
    "ensemble": {"params": "17.22M", "description": "SOTA Production Ensemble - Conservative Fusion of Base + Large with Heatmap Max-Gating"},
}

PREDICTION_CACHE: OrderedDict = OrderedDict()
MAX_CACHE_SIZE = 256

model_registry: Dict[str, TamperingModelServer] = {}
DEFAULT_VARIANT = os.environ.get("PATCH_MFFT_VARIANT", "base")
IMG_SIZE = int(os.environ.get("PATCH_MFFT_IMG_SIZE", "384"))


@asynccontextmanager
async def lifespan(app: FastAPI):
    print("[startup] Initializing Patch-MFFT model registry...")
    ckpt_dir_env = os.environ.get("PATCH_MFFT_CHECKPOINT_DIR")
    if ckpt_dir_env and Path(ckpt_dir_env).exists():
        ckpt_dir = Path(ckpt_dir_env)
    elif (Path(__file__).parent.parent / "checkpoints").exists():
        ckpt_dir = Path(__file__).parent.parent / "checkpoints"
    else:
        ckpt_dir = Path(__file__).parent.parent / "model" / "checkpoints"

    for variant in ["tiny", "base", "large"]:
        ckpt_path = ckpt_dir / f"{variant}.pt"
        try:
            model_registry[variant] = TamperingModelServer(
                str(ckpt_path) if ckpt_path.exists() else None,
                variant=variant,
                img_size=IMG_SIZE,
            )
        except Exception as e:
            print(f"[startup] Failed to initialize variant '{variant}': {e}")
    print(f"[startup] Loaded variants: {list(model_registry.keys())}")
    yield
    model_registry.clear()


app = FastAPI(
    title="Patch-MFFT Tampering Forensics API",
    description="Multi-Frequency Transformer for Localized Inpainting, Two-Person Splicing, and Composite Blackmail Detection",
    version="1.0.0",
    lifespan=lifespan,
)

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


@app.get("/", response_model=HealthResponse)
async def health():
    return HealthResponse(
        status="healthy",
        version="1.0.0",
        models_loaded=list(model_registry.keys()),
        timestamp=datetime.now(timezone.utc).isoformat(),
    )


@app.get("/models", response_model=ModelsResponse)
async def list_models():
    default_model = "ensemble" if ("base" in model_registry and "large" in model_registry) else (
        DEFAULT_VARIANT if DEFAULT_VARIANT in model_registry else (
            list(model_registry.keys())[0] if model_registry else "base"
        )
    )
    return ModelsResponse(
        default=default_model,
        models=[
            ModelInfo(
                id=v,
                loaded=(v in model_registry) or (v == "ensemble" and "base" in model_registry and "large" in model_registry),
                params=MODEL_INFO[v]["params"],
                description=MODEL_INFO[v]["description"],
            )
            for v in MODEL_INFO
        ],
    )


@app.post("/predict", response_model=PredictionResponse)
async def predict(
    file: UploadFile = File(...),
    model: str = "ensemble",
):
    variant = (model or "ensemble").lower()
    valid_variants = list(model_registry.keys()) + (["ensemble"] if "base" in model_registry and "large" in model_registry else [])
    if variant not in valid_variants:
        raise HTTPException(
            status_code=404,
            detail=f"Model variant '{variant}' not found. Available: {valid_variants}",
        )

    if not file.content_type or not file.content_type.startswith("image/"):
        raise HTTPException(status_code=400, detail="Uploaded file must be a valid image.")

    content = await file.read()
    if len(content) > 25 * 1024 * 1024:
        raise HTTPException(status_code=400, detail="Image size exceeds 25 MB limit.")

    sha256_hash = hashlib.sha256(content).hexdigest()
    cache_key = f"{sha256_hash}_{variant}"

    # Instant In-Memory LRU Cache Check (< 1 ms latency)
    if cache_key in PREDICTION_CACHE:
        cached_result = dict(PREDICTION_CACHE[cache_key])
        cached_result["cached"] = True
        PREDICTION_CACHE.move_to_end(cache_key)
        return PredictionResponse(**cached_result)

    try:
        image = Image.open(io.BytesIO(content)).convert("RGB")
    except Exception:
        raise HTTPException(status_code=400, detail="Corrupted image file.")

    if variant == "ensemble":
        result = predict_ensemble(
            model_registry["base"],
            model_registry["large"],
            image,
            filename=file.filename,
            sha256_hash=sha256_hash,
        )
    else:
        server = model_registry[variant]
        result = server.predict(image, filename=file.filename, sha256_hash=sha256_hash)

    # Save to high-speed LRU memory cache
    PREDICTION_CACHE[cache_key] = result
    if len(PREDICTION_CACHE) > MAX_CACHE_SIZE:
        PREDICTION_CACHE.popitem(last=False)

    return PredictionResponse(**result)