"""
================================================================================
  VERIDEX — DeepFake Worker Space  (Generic Template)
  ─────────────────────────────────────────────────────
  DEPLOY INSTRUCTIONS  — zero code changes between workers
  ──────────────────────────────────────────────────────────
  1. Commit this IDENTICAL app.py to all 7 Worker Spaces.
  2. Upload each worker's .pt weight file to its Space's files tab.
  3. In each Space → Settings → Variables, set:

       WEIGHT_FILE  = final_111_DeepFakeClassifier_tf_efficientnet_b7_ns_0_36
       MODEL_CLASS  = base        # or srm / gwap  (optional, default: base)

  4. That's it. No code edits required.

  API CONTRACT (called by the Master UI)
  ───────────────────────────────────────
  Input  : a .npy file (uint8, shape [N, H, W, 3], HWC, 380×380)
  Output : JSON  { "predictions": [float, ...], "n_frames": int }
            OR   { "error": "...", "predictions": null }

  GRADIO VERSION NOTE
  ────────────────────
  HF Spaces force-installs gradio==6.x at build time regardless of what
  requirements.txt pins. This file targets Gradio 6:
    • gr.File input passes a tempfile.SpooledTemporaryFile-backed object
      with a .name attribute in Gradio 6 (not a plain string or dict).
    • allow_flagging is removed (deprecated in Gradio 6; raises a warning
      that can abort startup on strict HF runtime configs).
================================================================================
"""

import os
import io
import re
import traceback
import logging

import numpy as np
import torch
import torch.nn as nn
from torch.nn.modules.dropout import Dropout
from torch.nn.modules.linear import Linear
from torch.nn.modules.pooling import AdaptiveAvgPool2d
from torchvision.transforms import Normalize
from functools import partial

import gradio as gr

# ── timm / efficientnet ───────────────────────────────────────────────────────
try:
    from timm.models.efficientnet import tf_efficientnet_b7_ns
except ImportError:
    # timm ≥ 0.9 moved the alias; fall back gracefully
    import timm
    tf_efficientnet_b7_ns = partial(timm.create_model, "tf_efficientnet_b7.ns_jft_in1k")


logging.basicConfig(level=logging.INFO, format="%(asctime)s [WORKER] %(levelname)s %(message)s")
logger = logging.getLogger(__name__)

# ══════════════════════════════════════════════════════════════════════════════
#  ❶  ALL CONFIG IS VIA ENV VARS — set these in each Space's Settings → Variables
#     WEIGHT_FILE : filename of the .pt checkpoint (no extension required)
#     MODEL_CLASS : "base" | "srm" | "gwap"  (default: base)
#     MINI_BATCH  : frames per forward pass   (default: 8)
#     WEIGHTS_DIR : directory containing the .pt file (default: repo root ".")
# ══════════════════════════════════════════════════════════════════════════════
WEIGHT_FILE = os.environ.get(
    "WEIGHT_FILE",
    "final_555_DeepFakeClassifier_tf_efficientnet_b7_ns_0_19",  # safe default
)
MODEL_CLASS = os.environ.get("MODEL_CLASS", "base")   # "base" | "srm" | "gwap"
MINI_BATCH  = int(os.environ.get("MINI_BATCH", "8"))   # frames per forward pass
WEIGHTS_DIR = os.environ.get("WEIGHTS_DIR", ".")        # dir that contains the .pt
# ══════════════════════════════════════════════════════════════════════════════

# ── ImageNet normalisation ────────────────────────────────────────────────────
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD  = [0.229, 0.224, 0.225]
normalize_fn  = Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD)

# ── EfficientNet-B7 feature size ──────────────────────────────────────────────
ENCODER_FEATURES = 2560

# ─────────────────────────────────────────────────────────────────────────────
#  Model definitions  (identical to deepfake_det.py so checkpoints load clean)
# ─────────────────────────────────────────────────────────────────────────────

def _make_encoder():
    return tf_efficientnet_b7_ns(pretrained=False, drop_path_rate=0.2)


def _setup_srm_weights(input_channels: int = 3) -> torch.Tensor:
    srm_kernel = torch.from_numpy(np.array([
        [[0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.],[0.,1.,-2.,1.,0.],[0.,0.,0.,0.,0.],[0.,0.,0.,0.,0.]],
        [[0.,0.,0.,0.,0.],[0.,-1.,2.,-1.,0.],[0.,2.,-4.,2.,0.],[0.,-1.,2.,-1.,0.],[0.,0.,0.,0.,0.]],
        [[-1.,2.,-2.,2.,-1.],[2.,-6.,8.,-6.,2.],[-2.,8.,-12.,8.,-2.],[2.,-6.,8.,-6.,2.],[-1.,2.,-2.,2.,-1.]],
    ])).float()
    srm_kernel[0] /= 2
    srm_kernel[1] /= 4
    srm_kernel[2] /= 12
    return srm_kernel.view(3, 1, 5, 5).repeat(1, input_channels, 1, 1)


def _setup_srm_layer(input_channels: int = 3) -> nn.Module:
    weights = _setup_srm_weights(input_channels)
    conv = nn.Conv2d(input_channels, 3, kernel_size=5, stride=1, padding=2, bias=False)
    with torch.no_grad():
        conv.weight = nn.Parameter(weights, requires_grad=False)
    return conv


class DeepFakeClassifier(nn.Module):
    def __init__(self, dropout_rate=0.0):
        super().__init__()
        self.encoder  = _make_encoder()
        self.avg_pool = AdaptiveAvgPool2d((1, 1))
        self.dropout  = Dropout(dropout_rate)
        self.fc       = Linear(ENCODER_FEATURES, 1)

    def forward(self, x):
        x = self.encoder.forward_features(x)
        x = self.avg_pool(x).flatten(1)
        x = self.dropout(x)
        return self.fc(x)


class DeepFakeClassifierSRM(nn.Module):
    def __init__(self, dropout_rate=0.5):
        super().__init__()
        self.encoder  = _make_encoder()
        self.avg_pool = AdaptiveAvgPool2d((1, 1))
        self.srm_conv = _setup_srm_layer(3)
        self.dropout  = Dropout(dropout_rate)
        self.fc       = Linear(ENCODER_FEATURES, 1)

    def forward(self, x):
        noise = self.srm_conv(x)
        x = self.encoder.forward_features(noise)
        x = self.avg_pool(x).flatten(1)
        x = self.dropout(x)
        return self.fc(x)


class _GWAP(nn.Module):
    def __init__(self, features: int):
        super().__init__()
        self.conv = nn.Conv2d(features, 1, kernel_size=1, bias=True)

    def forward(self, x):
        w = self.conv(x).sigmoid().exp()
        w = w / w.sum(dim=[2, 3], keepdim=True)
        return (w * x).sum(dim=[2, 3], keepdim=False)


class DeepFakeClassifierGWAP(nn.Module):
    def __init__(self, dropout_rate=0.5):
        super().__init__()
        self.encoder  = _make_encoder()
        self.avg_pool = _GWAP(ENCODER_FEATURES)
        self.dropout  = Dropout(dropout_rate)
        self.fc       = Linear(ENCODER_FEATURES, 1)

    def forward(self, x):
        x = self.encoder.forward_features(x)
        x = self.avg_pool(x)
        x = self.dropout(x)
        return self.fc(x)


_MODEL_MAP = {
    "base": DeepFakeClassifier,
    "srm":  DeepFakeClassifierSRM,
    "gwap": DeepFakeClassifierGWAP,
}


# ─────────────────────────────────────────────────────────────────────────────
#  Model loading  (runs once at startup)
# ─────────────────────────────────────────────────────────────────────────────

def load_model() -> nn.Module:
    device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    cls     = _MODEL_MAP.get(MODEL_CLASS, DeepFakeClassifier)
    model   = cls().to(device)

    weight_path = os.path.join(WEIGHTS_DIR, WEIGHT_FILE)
    # Allow common extensions in case the file was renamed
    if not os.path.exists(weight_path):
        for ext in (".pt", ".pth", ".bin"):
            if os.path.exists(weight_path + ext):
                weight_path = weight_path + ext
                break

    if not os.path.exists(weight_path):
        raise FileNotFoundError(
            f"Weight file not found: {weight_path}\n"
            f"Files present in '{WEIGHTS_DIR}': {os.listdir(WEIGHTS_DIR)}"
        )

    logger.info(f"Loading weights from: {weight_path}")

    # PyTorch 2.6+ requires weights_only=False for pickled checkpoints; also
    # use map_location='cpu' so the model loads on any machine regardless of
    # how it was saved.
    checkpoint  = torch.load(weight_path, map_location="cpu", weights_only=False)
    state_dict  = checkpoint.get("state_dict", checkpoint)

    # Strip "module." prefix added by DataParallel / DistributedDataParallel
    cleaned = {re.sub(r"^module\.", "", k): v for k, v in state_dict.items()}
    model.load_state_dict(cleaned, strict=True)
    model.eval()

    # FP16 halves VRAM; safe on both CUDA and CPU
    model = model.half()

    logger.info(f"Model ready — class={cls.__name__}, device={device}, fp16=True")
    return model, device


try:
    MODEL, DEVICE = load_model()
    LOAD_ERROR    = None
except Exception as exc:
    MODEL      = None
    DEVICE     = None
    LOAD_ERROR = traceback.format_exc()
    logger.error(f"MODEL LOAD FAILED:\n{LOAD_ERROR}")


# ─────────────────────────────────────────────────────────────────────────────
#  Inference helper
# ─────────────────────────────────────────────────────────────────────────────

def _preprocess_npy(npy_input) -> torch.Tensor:
    """
    Load a uint8 HWC .npy face-batch, convert to normalised float CHW tensor.

    Gradio version compatibility matrix
    ─────────────────────────────────────
    Gradio 4 : passes a plain string filepath  "/tmp/gradio/.../faces.npy"
    Gradio 4 : may wrap in dict {"path": "...", "orig_name": "..."}
    Gradio 6 : passes a tempfile.SpooledTemporaryFile (file-like with .name)
               OR a gradio.FileData dataclass with a .path attribute

    We resolve all four forms to a final file path or file-like object
    that np.load() can consume.
    """
    npy_path = None   # will hold a string path if resolvable
    file_obj = None   # will hold a file-like if path is unavailable

    # ── Form 1: plain string ──────────────────────────────────────────────────
    if isinstance(npy_input, str):
        npy_path = npy_input

    # ── Form 2: Gradio 4 dict {"path": ..., "orig_name": ...} ────────────────
    elif isinstance(npy_input, dict):
        npy_path = (
            npy_input.get("path")
            or npy_input.get("name")
            or next(iter(npy_input.values()), None)
        )

    # ── Form 3: Gradio 6 dataclass (has .path attribute) ─────────────────────
    elif hasattr(npy_input, "path"):
        npy_path = npy_input.path

    # ── Form 4: file-like object (SpooledTemporaryFile, BytesIO, etc.) ────────
    elif hasattr(npy_input, "read"):
        # Try to get the backing file path first (avoids reading into RAM twice)
        backing = getattr(npy_input, "name", None)
        if backing and isinstance(backing, str) and os.path.exists(backing):
            npy_path = backing
        else:
            file_obj = npy_input

    else:
        raise TypeError(
            f"Cannot resolve npy input of type {type(npy_input)}: {npy_input!r}"
        )

    # ── Load the array ─────────────────────────────────────────────────────────
    def _load(src):
        try:
            return np.load(src, allow_pickle=False)
        except ValueError:
            # Legacy pickled .npy — seek back to start if file-like
            if hasattr(src, "seek"):
                src.seek(0)
            return np.load(src, allow_pickle=True)

    if npy_path is not None:
        if not os.path.exists(npy_path):
            raise FileNotFoundError(f"NPY payload not found at: {npy_path}")
        faces_uint8 = _load(npy_path)
    else:
        faces_uint8 = _load(file_obj)

    # ── Validate shape ─────────────────────────────────────────────────────────
    if faces_uint8.ndim != 4 or faces_uint8.shape[3] != 3:
        raise ValueError(
            f"Expected uint8 array shape (N, H, W, 3), got {faces_uint8.shape}"
        )

    # Convert: uint8 HWC  →  float32 CHW  →  normalised
    tensor = torch.from_numpy(faces_uint8).float()   # [N, H, W, 3]
    tensor = tensor.permute(0, 3, 1, 2)              # [N, 3, H, W]
    # Normalise each frame in-place
    for i in range(tensor.shape[0]):
        tensor[i] = normalize_fn(tensor[i] / 255.0)

    return tensor  # float32, shape [N, 3, H, W]


def run_inference(tensor: torch.Tensor) -> list:
    """
    Forward-pass the pre-processed face tensor through the model in
    mini-batches of size MINI_BATCH to avoid OOM on 16 GB RAM spaces.
    Returns a flat Python list of per-frame fake-probabilities [0, 1].
    """
    predictions = []
    n = tensor.shape[0]

    with torch.no_grad():
        for start in range(0, n, MINI_BATCH):
            batch = tensor[start : start + MINI_BATCH]
            batch = batch.to(DEVICE).half()     # fp16 matches model dtype

            logits = MODEL(batch)               # [B, 1]
            probs  = torch.sigmoid(logits.squeeze(-1))  # [B]
            predictions.extend(probs.cpu().float().tolist())

    return predictions


# ─────────────────────────────────────────────────────────────────────────────
#  Gradio endpoint  (headless — no UI blocks, purely an API)
# ─────────────────────────────────────────────────────────────────────────────

def predict(npy_file) -> dict:
    """
    Gradio API endpoint.

    Parameters
    ----------
    npy_file : str | dict
        Filepath (or Gradio file dict) pointing to the .npy face batch.

    Returns
    -------
    dict with keys:
        predictions : list[float] | None
        n_frames    : int
        error       : str | None
    """
    if MODEL is None:
        msg = f"Model failed to load at startup:\n{LOAD_ERROR}"
        logger.error(msg)
        return {"predictions": None, "n_frames": 0, "error": msg}

    try:
        tensor      = _preprocess_npy(npy_file)
        n_frames    = tensor.shape[0]
        predictions = run_inference(tensor)
        logger.info(f"Inference OK — frames={n_frames}, mean_pred={np.mean(predictions):.4f}")
        return {"predictions": predictions, "n_frames": n_frames, "error": None}

    except Exception:
        err = traceback.format_exc()
        logger.error(f"Inference failed:\n{err}")
        return {"predictions": None, "n_frames": 0, "error": err}


# ─────────────────────────────────────────────────────────────────────────────
#  Launch
# ─────────────────────────────────────────────────────────────────────────────

demo = gr.Interface(
    fn=predict,
    inputs=gr.File(label="Face batch (.npy)", file_types=[".npy"]),
    outputs=gr.JSON(label="Worker prediction"),
    title=f"VERIDEX Worker — {WEIGHT_FILE}",
    description=(
        "Headless inference worker. "
        "POST a uint8 .npy face-batch; receive per-frame fake probabilities."
    ),
    # allow_flagging removed: deprecated in Gradio 5, gone in Gradio 6
)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860, show_error=True)