"""VideoSeal/PixelSeal watermark detection service.

Detects Meta VideoSeal watermarks embedded in images and videos using the
videoseal_base model. Returns a probability score indicating likelihood
that the media contains a watermark (indicating AI generation).

Unlike detection models that analyze visual artifacts, this service
inspects only for embedded VideoSeal watermarks.
"""

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

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

# ── Logging ────────────────────────────────────────────────────────────────

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

# ── Config ─────────────────────────────────────────────────────────────────

MODEL_NAME = "videoseal_detector"
MODEL_PORT = int(os.environ.get("MODEL_PORT", "9004"))
MAX_IMAGE_DIMENSION = 4096

_PRODUCTION = os.environ.get("PRODUCTION", "false").lower() == "true"

DETECTION_THRESHOLD = 0.6  # Bit confidence ratio: clean ~0.01, watermarked ~0.76
MAX_FRAMES = 16


def _get_device():
    """Select optimal device: CUDA > MPS > 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 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")


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_mem / 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


# ── Pydantic models ───────────────────────────────────────────────────────


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

    image_data: Optional[str] = None
    video_data: Optional[str] = None
    threshold: float = 0.5


# ── Confidence mapping ────────────────────────────────────────────────────


def _map_confidence_to_probability(confidence: float) -> float:
    """Map VideoSeal detection confidence to a probability score.

    Confidence values below the detection threshold are treated as
    noise and mapped to 0.5 (neutral). Values at or above the
    threshold are linearly scaled into [0.5, 1.0].

    Args:
        confidence: Raw detection confidence from VideoSeal in [0, 1].

    Returns:
        Probability in [0.5, 1.0]. 0.5 means neutral (no watermark).
    """
    if confidence < DETECTION_THRESHOLD:
        return 0.5
    return 0.5 + confidence * 0.5


# ── Model loading ─────────────────────────────────────────────────────────


def load_model_internal():
    """Load VideoSeal detector model 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 VideoSeal detector model...")
        try:
            # Use absolute path to model card (videoseal uses relative
            # Path("videoseal/cards") which fails when CWD != package root)
            import pathlib

            import videoseal

            cards_dir = pathlib.Path(videoseal.__file__).parent / "cards"
            card_path = cards_dir / "videoseal_1.0.yaml"
            loaded = videoseal.load(card_path)
            loaded = loaded.to(DEVICE)
            loaded.eval()
            model = loaded
            last_used_time = time.time()
            logger.info("VideoSeal detector ready on %s.", str(DEVICE))

        except Exception as exc:
            logger.exception("Failed to load VideoSeal detector: %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 VideoSeal detector to free RAM.")
            del model
            model = None
            gc.collect()


# ── Detection helpers ─────────────────────────────────────────────────────


def _detect_image(image_bytes: bytes) -> float:
    """Detect VideoSeal watermark in image bytes.

    Decodes the image, converts to a [1, 3, H, W] tensor in [0, 1],
    and runs the VideoSeal detector. Returns the mean p-value as the
    raw confidence score.

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

    Returns:
        Confidence float in [0, 1].
    """
    try:
        pil_image = Image.open(io.BytesIO(image_bytes))
        if (
            pil_image.width > MAX_IMAGE_DIMENSION
            or pil_image.height > MAX_IMAGE_DIMENSION
        ):
            logger.warning(
                "Image too large: %dx%d, max %d",
                pil_image.width,
                pil_image.height,
                MAX_IMAGE_DIMENSION,
            )
            return 0.0
        pil_image = pil_image.convert("RGB")
    except Exception as exc:
        logger.warning("Failed to decode image: %s", exc)
        return 0.0

    tensor = TF.to_tensor(pil_image).unsqueeze(0).to(DEVICE)  # [1, 3, H, W]

    with torch.no_grad():
        outputs = model.detect(tensor)

    # VideoSeal outputs 'preds': [1, nbits+1]. The first value is a
    # detection logit (unreliable). The real signal is the message bits:
    # watermarked images have many confident bits (far from 0.5),
    # clean images have bits clustered near 0.5.
    preds = outputs["preds"]
    bit_confs = torch.sigmoid(preds[0, 1:])  # message bit confidences
    confident_bits = ((bit_confs > 0.8) | (bit_confs < 0.2)).float().mean()
    confidence = float(confident_bits)
    return confidence


def _detect_video(video_bytes: bytes) -> float:
    """Detect VideoSeal watermark in video bytes.

    Writes bytes to a temporary file, extracts up to MAX_FRAMES
    evenly sampled frames via OpenCV, runs detection on each frame,
    and returns the maximum confidence across all frames.

    Args:
        video_bytes: Raw bytes of a video file.

    Returns:
        Maximum confidence float across frames in [0, 1].
    """
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
            tmp.write(video_bytes)
            tmp_path = tmp.name

        cap = cv2.VideoCapture(tmp_path)
        if not cap.isOpened():
            logger.warning("Failed to open video file.")
            return 0.0

        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
        if total_frames <= 0:
            cap.release()
            return 0.0

        # Sample frame indices evenly
        if total_frames <= MAX_FRAMES:
            indices = list(range(total_frames))
        else:
            indices = np.linspace(0, total_frames - 1, MAX_FRAMES, dtype=int).tolist()

        max_confidence = 0.0

        for idx in indices:
            cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
            ret, frame = cap.read()
            if not ret or frame is None:
                continue

            # OpenCV BGR -> RGB -> PIL -> tensor
            rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            pil_frame = Image.fromarray(rgb_frame)
            tensor = TF.to_tensor(pil_frame).unsqueeze(0).to(DEVICE)

            with torch.no_grad():
                outputs = model.detect(tensor)

            preds = outputs["preds"]
            bit_confs = torch.sigmoid(preds[0, 1:])
            confidence = float(((bit_confs > 0.8) | (bit_confs < 0.2)).float().mean())
            if confidence > max_confidence:
                max_confidence = confidence

        cap.release()
        return max_confidence

    except Exception as exc:
        logger.warning("Video detection error: %s", exc)
        return 0.0

    finally:
        if tmp_path and os.path.exists(tmp_path):
            try:
                os.unlink(tmp_path)
            except OSError:
                pass


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

app = FastAPI(
    title="VideoSeal Watermark Detection Service",
    description=(
        "Detects Meta VideoSeal/PixelSeal watermarks embedded in images "
        "and videos using the videoseal_base model."
    ),
    version="1.0.0",
    docs_url=None if _PRODUCTION else "/docs",
    redoc_url=None if _PRODUCTION else "/redoc",
    openapi_url=None if _PRODUCTION else "/openapi.json",
)

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


@app.get("/")
async def root():
    """Root endpoint with service information."""
    if _PRODUCTION:
        return {"status": "ok"}
    return {
        "model_name": MODEL_NAME,
        "description": (
            "VideoSeal/PixelSeal watermark detector for "
            "AI-generated images and videos"
        ),
        "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_mem / 1024**2
            ),
        }
    return {}


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


@app.post("/predict")
async def predict(payload: MediaInput) -> Dict[str, Any]:
    """Detect VideoSeal watermark in base64-encoded image or video.

    Accepts image_data or video_data as base64-encoded file bytes.
    Determines modality from which key is non-None. Returns a
    standard detection response with probability based on watermark
    evidence.

    Args:
        payload: Base64-encoded media and optional threshold.

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

        start = time.time()

        # Determine modality from input keys
        if payload.image_data is not None:
            raw_b64 = payload.image_data
            modality = "image"
        elif payload.video_data is not None:
            raw_b64 = payload.video_data
            modality = "video"
        else:
            raise HTTPException(
                status_code=400,
                detail=(
                    "No media data provided. Include one of: " "image_data, video_data."
                ),
            )

        # Decode base64 payload
        try:
            media_bytes = base64.b64decode(raw_b64)
        except Exception as exc:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid base64 data: {exc}",
            )

        if not media_bytes:
            raise HTTPException(
                status_code=400,
                detail="Empty media data after base64 decode.",
            )

        # Run detection based on modality
        if modality == "image":
            confidence = _detect_image(media_bytes)
        else:
            confidence = _detect_video(media_bytes)

        probability = _map_confidence_to_probability(confidence)
        prediction = 1 if probability > payload.threshold else 0
        class_label = "fake" if prediction == 1 else "real"
        inference_time = time.time() - start

        logger.info(
            "VideoSeal check (%s): %s (conf=%.4f, prob=%.4f, %.3fs)",
            modality,
            class_label,
            confidence,
            probability,
            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),
            "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():
    """Startup: preload model if configured, else lazy-load on first request."""
    if PRELOAD_MODEL:
        logger.info("Preloading VideoSeal detector at startup.")
        try:
            load_model_internal()
        except Exception as exc:
            logger.error("Preload failed: %s", exc)
    else:
        logger.info("VideoSeal 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__":
    logger.info("Starting VideoSeal detector service on port %d", MODEL_PORT)
    uvicorn.run(app, host="0.0.0.0", port=MODEL_PORT)