"""
CLIP-based binary classifier.

Strategy
--------
Use a frozen pretrained CLIP image encoder to produce a 512-d feature vector,
then a small linear head maps it to [P(authentic), P(ai_generated)].

Why CLIP?
~~~~~~~~~
CLIP's pretraining on 400M+ image-text pairs has been shown (UniversalFakeDetect,
Ojha et al. 2023, and follow-up work) to generalise far better across unseen
generators than ImageNet-pretrained CNNs. This is the single biggest lever for
cross-generator robustness in a commercial product.

Stage 1 status
--------------
The classification head is RANDOMLY INITIALIZED. This file establishes the
correct architecture and pipeline; predictions are not meaningful until the
head is trained on the curated dataset (see `scripts/dataset/`).

Licensing
---------
- HuggingFace `transformers` — Apache-2.0 (commercially safe)
- OpenAI CLIP ViT-B/32 weights — MIT (commercially safe)
"""
from __future__ import annotations

from pathlib import Path
from typing import Any

import torch
import torch.nn as nn
from PIL import Image

from ..config import settings
from .base import Detector, DetectorResult


class _LazyClipBackbone:
    """Defers loading transformers + downloading CLIP weights until first use.

    Keeps `import` cheap so tests that don't need the model don't pay the cost.
    """

    def __init__(self) -> None:
        self._model: Any | None = None
        self._processor: Any | None = None

    def get(self) -> tuple[Any, Any]:
        if self._model is None or self._processor is None:
            # Local import — transformers is heavy.
            from transformers import CLIPModel, CLIPProcessor

            name = settings.clip_model_name
            self._processor = CLIPProcessor.from_pretrained(name)
            self._model = CLIPModel.from_pretrained(name)
            self._model.eval()
            for p in self._model.parameters():
                p.requires_grad = False
        return self._model, self._processor


_clip_singleton = _LazyClipBackbone()


class ClipClassifier(Detector):
    """Frozen CLIP image encoder + small trainable head."""

    name = "clip_classifier"

    # CLIP ViT-B/32 image-embedding dim is 512.
    _EMBED_DIM = 512

    def __init__(self) -> None:
        # The trainable head. Two outputs: logits for [authentic, ai_generated].
        # Architecture mirrored exactly in scripts/train_head.py — keep them
        # in sync or checkpoint loads will fail.
        self._head: nn.Module = nn.Sequential(
            nn.Linear(self._EMBED_DIM, 256),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(256, 2),
        )
        self._head.eval()
        self._trained: bool = False

        # Stage 2 head loading. Priority:
        #   1. Local file at settings.head_checkpoint_path (dev / pre-baked).
        #   2. HF Hub model repo (settings.head_checkpoint_hf_repo).
        #   3. Fall through to scaffold mode (random head) — what tests use.
        # Errors here are caught and logged so the container always starts.
        ckpt = Path(settings.head_checkpoint_path)
        if not ckpt.is_file() and settings.head_checkpoint_hf_repo:
            try:
                from huggingface_hub import hf_hub_download

                downloaded = hf_hub_download(
                    repo_id=settings.head_checkpoint_hf_repo,
                    filename=settings.head_checkpoint_hf_filename,
                )
                ckpt = Path(downloaded)
            except Exception as exc:  # noqa: BLE001 — never crash on startup
                print(
                    f"WARNING: Stage 2 head download failed ({type(exc).__name__}: "
                    f"{exc}). Falling back to scaffold mode.",
                    flush=True,
                )

        if ckpt.is_file():
            self.load_head_weights(str(ckpt))

    @torch.no_grad()
    def run(self, image: Image.Image) -> DetectorResult:
        model, processor = _clip_singleton.get()

        # Processor handles resize, center-crop, normalize — same as CLIP's
        # original training preprocessing.
        inputs = processor(images=image, return_tensors="pt")
        features = model.get_image_features(**inputs)  # [1, 512]
        features = features / features.norm(p=2, dim=-1, keepdim=True)

        logits = self._head(features)  # [1, 2]
        probs = torch.softmax(logits, dim=-1)
        p_ai = float(probs[0, 1].item())

        return DetectorResult(
            name=self.name,
            score=p_ai,
            # Stage 2: trained on Open Images V7 (real) + Flux.1-schnell (AI).
            # Only authentic vs. ai_generated. Other classes left empty —
            # the ensemble will not assign mass to them.
            contributions={
                "authentic": 1.0 - p_ai,
                "ai_generated": p_ai,
            },
            notes=(
                None
                if self._trained
                else (
                    "Stage 1 scaffold — classification head is randomly initialised. "
                    "Output is not meaningful until Stage 2 fine-tuning."
                )
            ),
        )

    def load_head_weights(self, checkpoint_path: str) -> None:
        """Load fine-tuned head weights produced by Stage 2 training."""
        state = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
        self._head.load_state_dict(state)
        self._head.eval()
        self._trained = True