"""CO-SPY (CVPR 2025, Sony Research) AI-generated image detection service.

Wraps the CO-SPY dual-branch fusion detector with a FastAPI endpoint
matching the DeepSafe contract.  CO-SPY combines:
  - Semantic branch: frozen SigLIP ViT-SO400M-14-SigLIP-384 (open_clip)
    producing 1152-dim features -> Linear(1152, 1)
  - Artifact branch: SD v1-4 VAE encode+decode reconstruction error fed
    through a ResNet encoder producing 512-dim features -> Linear(512, 1)
  - Fusion layer: Linear(2, 1) calibrating the two branch logits

Key design decisions:
- Model architecture code is loaded from the cloned CO-SPY repo via
  sys.path at startup.  The API layer in this file is fully self-contained.
- Preprocessing mirrors the fusion detector's test_transform:
  Resize(384) -> CenterCrop(384) -> ToTensor(), with per-branch
  normalisation applied in the forward pass.
- Weights are downloaded from HuggingFace at model load time, not
  at build time.  Three weight files are needed: semantic_weights.pth,
  artifact_weights.pth, fusion_weights.pth.
- The SD v1-4 VAE and SigLIP model are loaded lazily by the detector
  constructors (heavy; ~2-3 GB combined).
"""

import base64
import gc
import io
import logging
import os
import platform
import sys
import threading
import time
from typing import Any, Dict, Optional

import numpy as np
import torch
import uvicorn
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from PIL import Image, ImageFile
from pydantic import BaseModel
from torchvision import transforms

ImageFile.LOAD_TRUNCATED_IMAGES = True

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    handlers=[logging.StreamHandler(sys.stdout)],
)
logger = logging.getLogger(__name__)

# -- Path setup ---------------------------------------------------------------
current_dir = os.path.dirname(os.path.abspath(__file__))
model_code_dir = os.path.join(current_dir, "model_code")

# Add the cloned CO-SPY repo root so that ``detectors``, ``utils``, and
# ``dataSets`` resolve during import.
if model_code_dir not in sys.path:
    sys.path.insert(0, model_code_dir)

# -- Config -------------------------------------------------------------------
MODEL_NAME = "cospy_detection"
WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", os.path.join(current_dir, "weights"))
HF_HOME = os.environ.get("HF_HOME", os.path.join(current_dir, "hf_cache"))
os.environ["HF_HOME"] = HF_HOME

# HuggingFace repo containing the three weight files.
HF_WEIGHTS_REPO = "ruojiruoli/Co-Spy-Pretrained-Weights"
HF_WEIGHTS_SUBDIR = "sd-v1_4"

# Expected checkpoint filenames inside WEIGHTS_DIR/<subdir>/
SEMANTIC_CKPT = "semantic_weights.pth"
ARTIFACT_CKPT = "artifact_weights.pth"
FUSION_CKPT = "fusion_weights.pth"

# Fusion preprocessing constants (matches fusion test_transform)
LOAD_SIZE = 384
CROP_SIZE = 384


def _get_device():
    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
    if override == "cpu":
        return torch.device("cpu")
    if override == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")
    if (
        override == "mps"
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        return torch.device("mps")
    if override:
        pass  # Invalid override, fall through to auto-detect
    if (
        platform.system() == "Darwin"
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


DEVICE = _get_device()
if DEVICE.type == "cuda":
    torch.backends.cudnn.benchmark = True
    torch.set_float32_matmul_precision("high")

if DEVICE.type == "cuda":
    logger.info(
        "Device: cuda (%s, %.1f GB VRAM)",
        torch.cuda.get_device_name(0),
        torch.cuda.get_device_properties(0).total_memory / 1024**3,
    )
else:
    logger.warning(
        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
        DEVICE,
    )

PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true"
MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", "600"))

# -- Globals ------------------------------------------------------------------
model = None
model_lock = threading.Lock()
last_used_time = 0


class ImageInput(BaseModel):
    """Request body for /predict."""

    image_data: str
    threshold: Optional[float] = 0.5


# -- Weight download ----------------------------------------------------------


def _ensure_weights_downloaded() -> str:
    """Download CO-SPY pretrained weights from HuggingFace if missing.

    Returns:
        Path to the directory containing the three weight files.
    """
    weights_subdir = os.path.join(WEIGHTS_DIR, HF_WEIGHTS_SUBDIR)
    semantic_path = os.path.join(weights_subdir, SEMANTIC_CKPT)
    artifact_path = os.path.join(weights_subdir, ARTIFACT_CKPT)
    fusion_path = os.path.join(weights_subdir, FUSION_CKPT)

    if (
        os.path.exists(semantic_path)
        and os.path.exists(artifact_path)
        and os.path.exists(fusion_path)
    ):
        logger.info("All CO-SPY weight files found in %s", weights_subdir)
        return weights_subdir

    logger.info(
        "Downloading CO-SPY weights from HuggingFace (%s)...",
        HF_WEIGHTS_REPO,
    )
    os.makedirs(weights_subdir, exist_ok=True)

    try:
        from huggingface_hub import hf_hub_download
    except ImportError:
        raise RuntimeError(
            "huggingface_hub is required to download weights. "
            "Install with: pip install huggingface_hub"
        )

    for fname in (SEMANTIC_CKPT, ARTIFACT_CKPT, FUSION_CKPT):
        remote_path = f"{HF_WEIGHTS_SUBDIR}/{fname}"
        dest = os.path.join(weights_subdir, fname)
        if not os.path.exists(dest):
            logger.info("  Downloading %s ...", remote_path)
            downloaded = hf_hub_download(
                repo_id=HF_WEIGHTS_REPO,
                filename=remote_path,
                local_dir=WEIGHTS_DIR,
            )
            logger.info("  Saved to %s", downloaded)

    return weights_subdir


# -- Preprocessing ------------------------------------------------------------


def preprocess_image(image_bytes: bytes) -> torch.Tensor:
    """Preprocess raw image bytes into a CO-SPY fusion-compatible tensor.

    Mirrors the fusion detector's test_transform:
    Resize(384) -> CenterCrop(384) -> ToTensor().  Per-branch normalisation
    (semantic: [0.5,0.5,0.5]/[0.5,0.5,0.5]; artifact: resize to 224 +
    identity normalisation) is applied inside the model's forward() method.

    Args:
        image_bytes: Raw bytes of a JPEG/PNG/etc. image.

    Returns:
        Tensor of shape [1, 3, 384, 384].

    Raises:
        Exception: If bytes cannot be decoded or processed.
    """
    pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

    transform = transforms.Compose(
        [
            transforms.Resize(LOAD_SIZE),
            transforms.CenterCrop(CROP_SIZE),
            transforms.ToTensor(),
        ]
    )

    tensor = transform(pil_image)  # [3, 384, 384]
    return tensor.unsqueeze(0)  # [1, 3, 384, 384]


# -- Model loading ------------------------------------------------------------


def load_model_internal():
    """Load the CO-SPY fusion detector onto the selected device."""
    global model, last_used_time

    with model_lock:
        if model is not None:
            last_used_time = time.time()
            return

        logger.info("Loading CO-SPY fusion model...")
        try:
            # Ensure weights are present (download from HF if needed)
            weights_subdir = _ensure_weights_downloaded()

            semantic_path = os.path.join(weights_subdir, SEMANTIC_CKPT)
            artifact_path = os.path.join(weights_subdir, ARTIFACT_CKPT)
            fusion_path = os.path.join(weights_subdir, FUSION_CKPT)

            # Import from the cloned CO-SPY repo.  The directory name
            # ``sd-v1_4`` contains a hyphen, so we must use importlib.
            import importlib

            detector_module = importlib.import_module("detectors.sd-v1_4")
            CoSpyFusionDetector = detector_module.CoSpyFusionDetector

            cospy_model = CoSpyFusionDetector(
                semantic_weights_path=semantic_path,
                artifact_weights_path=artifact_path,
            )
            cospy_model.load_weights(fusion_path)
            cospy_model.to(DEVICE)
            cospy_model.train(mode=False)

            model = cospy_model
            last_used_time = time.time()
            logger.info("CO-SPY fusion model ready on %s.", DEVICE)

        except Exception as exc:
            logger.exception("Failed to load CO-SPY model: %s", exc)
            model = None
            raise
        finally:
            gc.collect()


def ensure_model_loaded():
    """Load model on first request (lazy loading)."""
    global last_used_time
    if model is None:
        load_model_internal()
    else:
        last_used_time = time.time()


def unload_model_if_idle():
    """Evict model from RAM after MODEL_TIMEOUT seconds of inactivity."""
    global model
    if model is None or PRELOAD_MODEL:
        return
    with model_lock:
        if model is not None and (time.time() - last_used_time > MODEL_TIMEOUT):
            logger.info("Unloading idle CO-SPY model to free RAM.")
            del model
            model = None
            gc.collect()


# -- FastAPI app --------------------------------------------------------------

app = FastAPI(
    title="CO-SPY Deepfake Detection Service",
    description=(
        "AI-generated image detection using CO-SPY (CVPR 2025) -- "
        "dual-branch fusion of SigLIP semantic features and SD VAE "
        "reconstruction artifact features."
    ),
    version="1.0.0",
)

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


@app.get("/")
async def root():
    """Root endpoint with service information."""
    return {
        "model_name": MODEL_NAME,
        "description": ("CO-SPY (CVPR 2025) AI-generated image detector"),
        "device": str(DEVICE),
        "model_loaded": model is not None,
    }


def _gpu_health_info() -> dict:
    """Return GPU metrics for the health endpoint."""
    if torch.cuda.is_available() and DEVICE.type == "cuda":
        return {
            "gpu_name": torch.cuda.get_device_name(0),
            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
            "vram_total_mb": round(
                torch.cuda.get_device_properties(0).total_memory / 1024**2
            ),
        }
    return {}


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "model": MODEL_NAME,
        "device": str(DEVICE),
        "model_loaded": model is not None,
        **_gpu_health_info(),
    }


@app.post("/unload")
async def unload_model_endpoint():
    """Manually unload the model to free RAM."""
    global model
    if model is None:
        return {"status": "not_loaded"}
    del model
    model = None
    gc.collect()
    return {"status": "success", "message": "Model unloaded."}


@app.post("/predict")
async def predict(image_input: ImageInput) -> Dict[str, Any]:
    """Predict whether the submitted image is AI-generated.

    Args:
        image_input: Base64-encoded image and optional threshold.

    Returns:
        Dict with model name, fake probability, binary prediction,
        class label, and inference time.
    """
    try:
        ensure_model_loaded()
        if model is None:
            raise HTTPException(status_code=503, detail="Model not loaded.")

        start = time.time()

        try:
            image_bytes = base64.b64decode(image_input.image_data)
            input_tensor = preprocess_image(image_bytes)
            input_tensor = input_tensor.to(DEVICE)
        except Exception as exc:
            raise HTTPException(status_code=400, detail=f"Invalid image data: {exc}")

        with torch.no_grad():
            output = model(input_tensor)  # raw logit [B, 1]
            probability_fake = output.sigmoid().item()  # scalar in [0, 1]

        prediction = 1 if probability_fake >= image_input.threshold else 0
        class_label = "fake" if prediction == 1 else "real"
        inference_time = time.time() - start

        logger.info(
            "Prediction: %s (prob=%.4f, %.3fs)",
            class_label,
            probability_fake,
            inference_time,
        )

        if not PRELOAD_MODEL and MODEL_TIMEOUT > 0:
            threading.Timer(MODEL_TIMEOUT + 5.0, unload_model_if_idle).start()

        return {
            "model": MODEL_NAME,
            "probability": float(probability_fake),
            "prediction": int(prediction),
            "class": class_label,
            "inference_time": float(inference_time),
        }

    except HTTPException:
        raise
    except Exception as exc:
        logger.exception("Prediction error: %s", exc)
        raise HTTPException(status_code=500, detail=str(exc))


@app.on_event("startup")
async def startup_event():
    """Preload model if PRELOAD_MODEL=true, else lazy-load on first request."""
    if PRELOAD_MODEL:
        logger.info("Preloading CO-SPY model at startup.")
        try:
            load_model_internal()
        except Exception as exc:
            logger.error("Preload failed: %s", exc)
    else:
        logger.info("CO-SPY service ready -- model loads on first request.")

    if not PRELOAD_MODEL and MODEL_TIMEOUT > 0:

        def _periodic_check():
            unload_model_if_idle()
            threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start()

        threading.Timer(MODEL_TIMEOUT / 2.0, _periodic_check).start()


if __name__ == "__main__":
    port = int(os.environ.get("MODEL_PORT", 5007))
    logger.info("Starting CO-SPY service on port %d", port)
    uvicorn.run(app, host="0.0.0.0", port=port)