"""LipFD video deepfake detection service.

Wraps the LipFD (NeurIPS 2024) lip forgery detection model with a
FastAPI endpoint. Uses CLIP ViT-L/14 as a global feature extractor
and a ResNet-50-based region-aware classifier that operates on
multi-scale crops of detected faces.

Reference: Hou et al., "LipFD: Lip Forgery Detection via
Region-Aware CLIP", NeurIPS 2024.
"""

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 torchvision.transforms as transforms
import uvicorn
from facenet_pytorch import MTCNN
from fastapi import FastAPI, HTTPException
from PIL import Image
from pydantic import BaseModel, ConfigDict, Field

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

MODEL_PORT = int(os.environ.get("MODEL_PORT", 7006))
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/lipfd_checkpoint.pth")

# Add model_code to path so LipFD's internal imports resolve
_MODEL_CODE_DIR = os.environ.get(
    "MODEL_CODE_DIR",
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "model_code"),
)
if _MODEL_CODE_DIR not in sys.path:
    sys.path.insert(0, _MODEL_CODE_DIR)

# Number of frames to uniformly sample from the video
NUM_FRAMES = 32
# Face bounding-box margin factor (fraction of bbox dimension)
MARGIN_FACTOR = 0.5
# CLIP normalization constants (from OpenAI CLIP)
CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
# Multi-scale crop indices (from LipFD datasets.py)
# Scale 0 = full 224x224 crop, scale 1 = inner 168x168, scale 2 = inner 102x102
CROP_IDX = [(28, 196), (61, 163)]
# Number of spatial sub-crops per face (horizontal slices of lower-face)
NUM_SUBCROP_POSITIONS = 5


# ── Device selection ───────────────────────────────────────────────────────


def _get_device() -> torch.device:
    """Select optimal device: CUDA > MPS (Apple Silicon) > 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 = None
_face_detector: Optional[MTCNN] = None
_device: Optional[torch.device] = None
_load_lock = threading.Lock()


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

    if _model is not None:
        return

    with _load_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 LipFD model on %s ...", _device)

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

        # ── LipFD classifier ──────────────────────────────────────────
        from models import build_model  # noqa: E402

        model = build_model("CLIP:ViT-L/14")

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

        checkpoint = torch.load(WEIGHTS_PATH, map_location="cpu", weights_only=False)
        model.load_state_dict(checkpoint["model"])
        model = model.to(_device)
        model.eval()

        _model = model
        logger.info("LipFD 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="LipFD Detection Service",
    description=(
        "Lip Forgery Detection via Region-Aware CLIP " "(ViT-L/14, NeurIPS 2024)"
    ),
    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 = "lipfd_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": "lipfd_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": "lipfd_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 with a relative margin around the bbox.

    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_faces(frame: np.ndarray) -> List[np.ndarray]:
    """Detect faces in a frame and return margin-expanded crops.

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

    Returns:
        List of face-crop arrays (variable size, RGB uint8).
    """
    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
        crops.append(face)

    return crops


# ── LipFD preprocessing ───────────────────────────────────────────────────

_resize_224 = transforms.Resize((224, 224))
_resize_1120 = transforms.Resize((1120, 1120))
_clip_normalize = transforms.Normalize(mean=CLIP_MEAN, std=CLIP_STD)


def _prepare_lipfd_inputs(
    face_crop: np.ndarray,
) -> Tuple[torch.Tensor, List[List[torch.Tensor]]]:
    """Convert a single face crop into LipFD model inputs.

    Mirrors the preprocessing in LipFD's AVLip dataset:
      1. Convert face to float32 tensor (C, H, W), no /255 normalization
         (the original dataset uses raw cv2.imread pixel values).
      2. CLIP-normalize, then create 5 horizontal sub-crops of the
         lower-face region, each at 3 zoom scales (1.0x, 0.65x, 0.45x).
      3. Resize the raw image to 1120x1120 for the CLIP encoder.

    Args:
        face_crop: RGB uint8 numpy array of a detected face (H, W, 3).

    Returns:
        img_1120: Tensor of shape (3, 1120, 1120) for CLIP feature
            extraction (raw pixel values, no normalization).
        crops: List of 3 scale-lists, each containing 5 tensors of
            shape (3, 224, 224) -- CLIP-normalized.
    """
    # Raw float tensor (same as original: torch.tensor(cv2.imread(...)))
    # cv2.imread returns BGR but we already have RGB from our pipeline.
    # The original code reads BGR from cv2 and feeds it directly, so
    # we convert back to BGR to match training distribution.
    face_bgr = cv2.cvtColor(face_crop, cv2.COLOR_RGB2BGR)
    img_tensor = torch.tensor(face_bgr, dtype=torch.float32).permute(2, 0, 1)

    # CLIP-normalize for the crop pathway
    img_norm = _clip_normalize(img_tensor)

    # Build multi-scale crops from the lower-face region.
    # Original dataset: face images are 1000x(variable), crops come from
    # img[:, 500:, i:i+500] for 5 horizontal positions.
    # We adapt to arbitrary face dimensions: take the lower half and split
    # into 5 overlapping horizontal positions.
    _, h, w = img_norm.shape
    half_h = h // 2
    lower_face = img_norm[:, half_h:, :]  # lower half of the face

    lf_h, lf_w = lower_face.shape[1], lower_face.shape[2]
    # Create 5 square-ish sub-crops across the lower face width
    crop_size = min(lf_h, lf_w)
    if crop_size < 2:
        crop_size = max(lf_h, lf_w, 2)

    # Positions spread across the width
    if lf_w > crop_size:
        positions = np.linspace(0, lf_w - crop_size, NUM_SUBCROP_POSITIONS, dtype=int)
    else:
        positions = [0] * NUM_SUBCROP_POSITIONS

    # Scale 0 (1.0x): 5 crops from lower face, each resized to 224x224
    crops: List[List[torch.Tensor]] = [[], [], []]
    for pos in positions:
        patch = lower_face[:, :crop_size, pos : pos + crop_size]
        crop_224 = _resize_224(patch)
        crops[0].append(crop_224)

        # Scale 1 (0.65x): inner crop of the 224x224
        crops[1].append(
            _resize_224(
                crop_224[
                    :,
                    CROP_IDX[0][0] : CROP_IDX[0][1],
                    CROP_IDX[0][0] : CROP_IDX[0][1],
                ]
            )
        )
        # Scale 2 (0.45x): tighter inner crop
        crops[2].append(
            _resize_224(
                crop_224[
                    :,
                    CROP_IDX[1][0] : CROP_IDX[1][1],
                    CROP_IDX[1][0] : CROP_IDX[1][1],
                ]
            )
        )

    # Full image resized to 1120x1120 for CLIP global feature extraction
    img_1120 = _resize_1120(img_tensor)

    return img_1120, crops


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


@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    """Run LipFD 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. For each face, create multi-scale crops and CLIP features.
      5. Run the region-aware classifier.
      6. Per-frame max probability, then average across frames.

    Returns probability=0.5 (undetermined) if no faces are detected.
    """
    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:
            face_crops = _detect_faces(frame)
            if not face_crops:
                continue

            frame_scores: List[float] = []

            for face_crop in face_crops:
                img_1120, crops = _prepare_lipfd_inputs(face_crop)

                # Batch dim = 1 for single-face inference
                img_batch = img_1120.unsqueeze(0).to(_device)

                # Build crop tensors: each scale-list element gets batch dim
                crops_batch = [
                    [c.unsqueeze(0).to(_device) for c in scale_crops]
                    for scale_crops in crops
                ]

                with torch.no_grad():
                    # Get CLIP global features
                    features = _model.get_features(img_batch).to(_device)
                    # Run region-aware classifier
                    pred_score, _, _ = _model(crops_batch, features)
                    prob = pred_score.sigmoid().item()

                frame_scores.append(prob)
                total_faces += 1

            if frame_scores:
                per_frame_max.append(max(frame_scores))

        # ── Aggregate ──────────────────────────────────────────────
        if per_frame_max:
            probability = float(np.mean(per_frame_max))
        else:
            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 LipFD 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)