"""RECCE deepfake video detection service.

Wraps the RECCE (CVPR 2022) reconstruction-classification face forgery
detection model with a FastAPI endpoint.  Uses an Xception encoder with
guided attention and graph reasoning, trained via DeepfakeBench on
FaceForensics++ (c40).

The checkpoint originates from DeepfakeBench v1.0.1 which wraps the
original RECCE architecture under a ``model.`` prefix and trains with
2-class (real/fake) output instead of the original repo's 1-class
sigmoid.  Weights are loaded by stripping the ``model.`` prefix and
instantiating ``Recce(num_classes=2)``.

Reference: Cao et al., "End-to-End Reconstruction-Classification
Learning for Face Forgery Detection", CVPR 2022.
"""

import base64
import gc
import logging
import os
import platform
import sys
import tempfile
import threading
import time
from typing import Any, Dict, List, Optional, Tuple

import cv2
import numpy as np
import torch
import torch.nn.functional as F
import uvicorn
from facenet_pytorch import MTCNN
from fastapi import FastAPI, HTTPException
from PIL import Image
from pydantic import BaseModel, ConfigDict, Field

# ── Model code import ──────────────────────────────────────────────────────
# The original RECCE model code is vendored under model_code/.  We add it
# to sys.path so that ``from model.network import Recce`` resolves.
_MODEL_CODE_DIR = os.path.join(os.path.dirname(__file__), "model_code")
if _MODEL_CODE_DIR not in sys.path:
    sys.path.insert(0, _MODEL_CODE_DIR)

# Patch timm's xception to skip pretrained-weight download.  We load our
# own checkpoint, so downloading ImageNet weights wastes bandwidth and
# fails in air-gapped Docker containers.
import timm.models  # noqa: E402

_original_xception = timm.models.xception


def _xception_no_pretrained(**kwargs):
    """Force pretrained=False to avoid downloading default weights."""
    kwargs["pretrained"] = False
    return _original_xception(**kwargs)


timm.models.xception = _xception_no_pretrained

from model.network import Recce  # noqa: E402

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

MODEL_PORT = int(os.environ.get("MODEL_PORT", 7007))
PRELOAD_MODEL = os.environ.get("PRELOAD_MODEL", "false").lower() == "true"
MODEL_TIMEOUT = int(os.environ.get("MODEL_TIMEOUT", 1800))
WEIGHTS_PATH = os.environ.get("WEIGHTS_PATH", "/app/weights/recce_checkpoint.pth")

# RECCE uses 299x299 face crops (from config/Recce.yml and inference.py)
IMAGE_SIZE = (299, 299)
# Number of frames to sample from each video
NUM_FRAMES = 32
# Face crop margin factor
MARGIN_FACTOR = 0.5


def _get_device() -> torch.device:
    """Select optimal device: CUDA > MPS (Apple) > 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 torch.cuda.is_available():
        return torch.device("cuda")
    if (
        platform.system() == "Darwin"
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        return torch.device("mps")
    return torch.device("cpu")


# ── Global state ───────────────────────────────────────────────────────────

_model: Optional[Recce] = None
_face_detector: Optional[MTCNN] = None
_device: Optional[torch.device] = None
_load_lock = threading.Lock()


def _load_models() -> None:
    """Load RECCE detector and MTCNN face detector (thread-safe)."""
    global _model, _face_detector, _device

    if _model is not None:
        return

    with _load_lock:
        # Double-check after acquiring lock
        if _model is not None:
            return

        _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,
            )
        logger.info("Loading RECCE model on %s ...", _device)

        # ── Face detector (MTCNN) ──────────────────────────────────────
        _face_detector = MTCNN(
            keep_all=True,
            device=_device,
            post_process=False,
        )

        # ── RECCE classifier ───────────────────────────────────────────
        # DeepfakeBench trains RECCE with 2-class output (real, fake)
        net = Recce(num_classes=2)

        if not os.path.exists(WEIGHTS_PATH):
            raise FileNotFoundError(f"RECCE weights not found at {WEIGHTS_PATH}")

        checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False)

        # DeepfakeBench wraps the RECCE model under a 'model.' prefix.
        # Strip it so the keys match the original Recce class.
        if any(k.startswith("model.") for k in checkpoint.keys()):
            state_dict = {
                k[len("model.") :]: v
                for k, v in checkpoint.items()
                if k.startswith("model.")
            }
        else:
            state_dict = checkpoint

        net.load_state_dict(state_dict)
        net = net.to(_device)
        net.eval()

        _model = net
        logger.info("RECCE model loaded successfully.")


def _is_model_loaded() -> bool:
    """Return True if both the classifier and face detector are loaded."""
    return _model is not None and _face_detector is not None


# ── FastAPI app ────────────────────────────────────────────────────────────

app = FastAPI(
    title="RECCE Detection Service",
    description=(
        "End-to-End Reconstruction-Classification Learning "
        "for Face Forgery Detection (Xception backbone, CVPR 2022)"
    ),
    version="1.0.0",
)


class PredictRequest(BaseModel):
    """Incoming prediction request."""

    video_data: str  # Base64-encoded video bytes
    threshold: float = 0.5


class PredictResponse(BaseModel):
    """Outgoing prediction result."""

    model_config = ConfigDict(populate_by_name=True)

    model: str = "recce_detection"
    probability: float
    prediction: int
    class_name: str = Field(..., alias="class")
    inference_time: float
    metadata: Dict[str, Any]


@app.on_event("startup")
async def startup_event():
    """Optionally preload model at startup."""
    if PRELOAD_MODEL:
        _load_models()


@app.get("/")
def root():
    """Service info endpoint."""
    return {
        "service": "recce_detection",
        "port": MODEL_PORT,
        "model_loaded": _is_model_loaded(),
        "device": str(_device) if _device else "unknown",
    }


def _gpu_health_info() -> dict:
    """Return GPU metrics for the health endpoint."""
    if torch.cuda.is_available() and _device is not None 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")
def health():
    """Health check endpoint."""
    return {
        "status": "healthy",
        "model": "recce_detection",
        "device": str(_device) if _device else "cpu",
        "model_loaded": _is_model_loaded(),
        "weights_exist": os.path.exists(WEIGHTS_PATH),
        **_gpu_health_info(),
    }


# ── Video / face utilities ─────────────────────────────────────────────────


def _extract_frames(video_path: str, num_frames: int = NUM_FRAMES) -> List[np.ndarray]:
    """Uniformly sample *num_frames* RGB frames from a video file.

    Args:
        video_path: Path to the video on disk.
        num_frames: Number of frames to extract.

    Returns:
        List of RGB uint8 numpy arrays (H, W, 3).
    """
    cap = cv2.VideoCapture(video_path)
    total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total <= 0:
        cap.release()
        return []

    indices = np.linspace(0, total - 1, num_frames, endpoint=True, dtype=int)
    frames: List[np.ndarray] = []

    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, int(idx))
        ret, frame = cap.read()
        if ret:
            frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))

    cap.release()
    return frames


def _crop_face(
    img: np.ndarray,
    bbox: Tuple[float, float, float, float],
    margin: float = MARGIN_FACTOR,
) -> np.ndarray:
    """Crop a face region from an image with a relative margin.

    Args:
        img: RGB image array (H, W, 3).
        bbox: (x0, y0, x1, y1) face bounding box.
        margin: Fraction of bbox dimension to add as padding.

    Returns:
        Cropped face region as a numpy array.
    """
    h_img, w_img = img.shape[:2]
    x0, y0, x1, y1 = bbox
    w = x1 - x0
    h = y1 - y0

    x0_new = max(0, int(x0 - w * margin / 2))
    x1_new = min(w_img, int(x1 + w * margin / 2) + 1)
    y0_new = max(0, int(y0 - h * margin / 2))
    y1_new = min(h_img, int(y1 + h * margin / 2) + 1)

    return img[y0_new:y1_new, x0_new:x1_new]


def _detect_and_crop_faces(frame: np.ndarray) -> List[np.ndarray]:
    """Detect faces in a single frame and return cropped+resized chips.

    Uses MTCNN for face detection, then crops with margin and resizes
    to IMAGE_SIZE (299x299).

    Args:
        frame: RGB image array (H, W, 3).

    Returns:
        List of face crops resized to IMAGE_SIZE, as uint8 arrays.
    """
    assert _face_detector is not None

    pil_img = Image.fromarray(frame)
    boxes, _ = _face_detector.detect(pil_img)

    if boxes is None or len(boxes) == 0:
        return []

    crops: List[np.ndarray] = []
    for box in boxes:
        x0, y0, x1, y1 = box.tolist()
        face = _crop_face(frame, (x0, y0, x1, y1))
        if face.size == 0:
            continue
        resized = cv2.resize(face, IMAGE_SIZE)
        crops.append(resized)

    return crops


def _preprocess_face(face_crop: np.ndarray) -> torch.Tensor:
    """Apply RECCE-specific preprocessing to a face crop.

    RECCE uses Normalize(mean=[0.5]*3, std=[0.5]*3) which maps
    [0, 255] uint8 to [-1, 1] float32.  This matches the
    albumentations pipeline in the original inference.py.

    Args:
        face_crop: RGB uint8 array of shape (299, 299, 3).

    Returns:
        Tensor of shape (3, 299, 299) in range [-1, 1].
    """
    # uint8 [0, 255] -> float32 [0, 1] -> normalized [-1, 1]
    tensor = torch.tensor(face_crop).permute(2, 0, 1).float().div(255.0)
    tensor = (tensor - 0.5) / 0.5
    return tensor


# ── Prediction endpoint ────────────────────────────────────────────────────


@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    """Run RECCE face-forgery detection on a base64-encoded video.

    Pipeline:
      1. Decode video and write to temp file.
      2. Extract uniformly-sampled frames.
      3. Detect and crop faces per frame (MTCNN).
      4. Preprocess each face crop (resize 299x299, normalize to [-1,1]).
      5. Classify each face crop with RECCE (2-class softmax).
      6. For each frame, take the max fake probability across faces.
      7. Average the per-frame max probabilities.

    If no faces are detected in any frame the service returns
    probability=0.5 (undetermined) rather than raising an error.
    """
    if not _is_model_loaded():
        _load_models()

    start_time = time.time()

    # ── Decode video ───────────────────────────────────────────────
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
        try:
            video_bytes = base64.b64decode(request.video_data)
            tmp.write(video_bytes)
            tmp_path = tmp.name
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Failed to decode video: {e}")

    try:
        # ── Extract frames ─────────────────────────────────────────
        frames = _extract_frames(tmp_path, NUM_FRAMES)
        if not frames:
            raise HTTPException(
                status_code=400,
                detail="Could not extract frames from video.",
            )

        # ── Detect faces and classify ──────────────────────────────
        per_frame_max: List[float] = []
        total_faces = 0

        for frame in frames:
            crops = _detect_and_crop_faces(frame)
            if not crops:
                continue

            # Preprocess and build batch tensor
            tensors = [_preprocess_face(c) for c in crops]
            batch_tensor = torch.stack(tensors).to(_device)

            with torch.no_grad():
                logits = _model(batch_tensor)
                # RECCE forward() calls squeeze() which drops the
                # batch dim when batch=1, producing shape (2,)
                # instead of (1, 2).  Always ensure 2D.
                if logits.dim() == 1:
                    logits = logits.unsqueeze(0)
                probs = F.softmax(logits, dim=1)[:, 1]  # fake prob

            frame_max = probs.max().cpu().item()
            per_frame_max.append(frame_max)
            total_faces += len(crops)

        # ── Aggregate ──────────────────────────────────────────────
        if per_frame_max:
            probability = float(np.mean(per_frame_max))
        else:
            # No faces detected in any frame -- undetermined
            probability = 0.5

        prediction = 1 if probability >= request.threshold else 0
        class_name = "fake" if prediction == 1 else "real"

        return PredictResponse(
            probability=probability,
            prediction=prediction,
            class_name=class_name,
            inference_time=time.time() - start_time,
            metadata={
                "frames_sampled": len(frames),
                "frames_with_faces": len(per_frame_max),
                "total_faces_detected": total_faces,
                "device": str(_device),
            },
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.exception("Error during RECCE prediction")
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        gc.collect()


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT)