"""Local NSFW model adapter.

Wraps the on-device ViT classifier behind a small interface so the safety service
depends on ``NsfwModel.score`` rather than on torch/transformers or the archived
CLI. Heavy imports are lazy: importing this module never loads the model, so the
rest of the app (and its tests) stay light. Scoring itself is non-deterministic
across hardware and is exercised in dedicated model tests, not the unit suite.

Extracted from nsfwtag.scoring.score_images (donor_ledger.yaml: nt-score-model).
The donor's CSV cache is dropped here — scores are persisted in the database by
asset ID (safety repository), not a path-keyed CSV.
"""

from __future__ import annotations

from pathlib import Path

from photo_pipeline import imaging

MODEL_ID = "AdamCodd/vit-base-nsfw-detector"
BATCH = 16


class NsfwModel:
    def __init__(self, model_id: str = MODEL_ID, batch: int = BATCH) -> None:
        self.model_id = model_id
        self.batch = batch
        self._model = None
        self._device = None
        self._size = None
        self._nsfw_idx = None

    def _ensure_loaded(self) -> None:
        if self._model is not None:
            return
        import torch
        from transformers import AutoModelForImageClassification

        self._device = "mps" if torch.backends.mps.is_available() else "cpu"
        try:
            model = AutoModelForImageClassification.from_pretrained(
                self.model_id, local_files_only=True
            )
        except Exception:
            model = AutoModelForImageClassification.from_pretrained(self.model_id)
        self._model = model.to(self._device).eval()
        self._nsfw_idx = next(
            i for i, label in model.config.id2label.items() if label.lower() == "nsfw"
        )
        self._size = getattr(model.config, "image_size", 224)

    def score(self, paths: list[Path | str]) -> list[tuple[str, float]]:
        """Return ``[(path, nsfw_probability)]`` for each readable image."""
        if not paths:
            return []
        self._ensure_loaded()
        import numpy as np
        import torch
        from PIL import Image

        # The donor set ``ImageFile.LOAD_TRUNCATED_IMAGES = True`` here. That flag is
        # process-global: in this application the same process also hashes files and
        # renders previews, and those must keep failing loudly on a truncated file
        # rather than quietly working on half of one (US07-03). An unreadable image
        # is skipped instead — it stays unscored, and therefore visibly undecided.
        def preprocess(path):
            with imaging.open_image(path) as image:
                small = image.convert("RGB").resize((self._size, self._size), Image.BILINEAR)
            array = (np.asarray(small, dtype="float32") / 255.0 - 0.5) / 0.5
            return torch.from_numpy(array).permute(2, 0, 1)

        results: list[tuple[str, float]] = []
        items = [str(p) for p in paths]
        for start in range(0, len(items), self.batch):
            tensors, batch_paths = [], []
            for path in items[start : start + self.batch]:
                try:
                    tensors.append(preprocess(path))
                    batch_paths.append(path)
                except (imaging.MediaError, OSError, ValueError):
                    # One bad file must not cost the batch its other fifteen.
                    continue
            if not tensors:
                continue
            with torch.no_grad():
                probs = self._model(
                    pixel_values=torch.stack(tensors).to(self._device)
                ).logits.softmax(-1)
            for path, prob in zip(batch_paths, probs):
                results.append((path, float(prob[self._nsfw_idx].item())))
        return results