"""Camie-tagger-v2 ONNX wrapper (CPU).

Single-image at a time. Loaded lazily inside the ml-worker process; NOT
thread-safe — the ml queue worker runs --concurrency=1 per process (scale ML by
running multiple worker replicas, not threads).

v2 layout reference: HuggingFace Camais03/camie-tagger-v2 root has
camie-tagger-v2.onnx (789 MB) + camie-tagger-v2-metadata.json (7.77 MB)
+ config.json. Tags ship as nested JSON, not CSV. Preprocessing and
output handling follow the published onnx_inference.py reference:
ImageNet normalize, NCHW layout, sigmoid on refined logits (output[1]).
"""

import json
import os
from dataclasses import dataclass
from pathlib import Path

import numpy as np
from PIL import Image, ImageFile

# Cap inference threads (see Tagger.load) so each ml-worker replica is a bounded
# core consumer on a shared node — keep N_replicas × this within the cores
# allotted to ML so replicas don't oversubscribe the box / starve the DB.
_INTRA_OP_THREADS = 4

# onnxruntime lives in requirements-ml.txt only — it is NOT installed in the
# lean web image or in CI. Imported lazily inside Tagger.load() so this module
# imports fine without it (the suggestion service imports SURFACED_CATEGORIES
# from here in the web container, and CI collects the pure-logic tests).

# Tolerate minutely-truncated source images (same rationale as IR's wd14.py:
# a few missing bytes at the JPEG EOI shouldn't block tagging the whole image).
ImageFile.LOAD_TRUNCATED_IMAGES = True

MODEL_NAME = os.environ.get("CAMIE_MODEL_NAME", "camie-tagger-v2")
_MODEL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "camie"
_MODEL_FILE = f"{MODEL_NAME}.onnx"
_METADATA_FILE = f"{MODEL_NAME}-metadata.json"

# Ingest floor below which predictions aren't stored (keeps the JSON compact).
# DEFAULT/fallback only — the live value is DB-backed
# (ml_settings.tagger_store_floor) and passed into infer() per call by the ml
# task. 0.70: the suggestion path already filters there and the centroid path
# covers lower-confidence preferred tags, so the sub-0.70 tail is redundant
# (it had bloated image_record's TOAST to ~100 GB; plan-task #764).
DEFAULT_STORE_FLOOR = 0.70

# The categories FC-2b surfaces in the UI. Others (meta/rating/year) are
# still stored but the suggestion service filters them out.
# 'artist' retired in FC-2d-vii-c — artist identity is acquisition-derived
# (image_record.artist_id), never ML-inferred. 'copyright' retired
# 2026-06-01 — operator doesn't use the copyright tag-kind; fandom is
# this app's franchise/series concept (per TagsView.vue's doc comment).
# Raw predictions for both categories still get stored at STORE_FLOOR but
# don't surface in suggestions.
SURFACED_CATEGORIES = {"character", "general"}

# ImageNet preprocessing constants (per Camie v2 onnx_inference.py).
_IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
# Square-pad color ≈ ImageNet mean × 255 (matches reference inference).
_PAD_COLOR = (124, 116, 104)


@dataclass(frozen=True)
class TagPrediction:
    name: str
    category: str
    confidence: float


class Tagger:
    def __init__(self, model_dir: Path | None = None):
        self._model_dir = model_dir or _MODEL_DIR
        self._session = None  # onnxruntime.InferenceSession once load()ed
        self._tag_names: list[str] | None = None
        self._tag_categories: list[str] | None = None
        self._input_name: str | None = None
        self._input_size: int = 512

    def load(self) -> None:
        if self._session is not None:
            return
        model_path = self._model_dir / _MODEL_FILE
        meta_path = self._model_dir / _METADATA_FILE
        if not model_path.is_file():
            raise RuntimeError(
                f"Camie {_MODEL_FILE} missing at {model_path}. "
                f"Populate /models via the ml-worker downloader."
            )
        if not meta_path.is_file():
            raise RuntimeError(
                f"Camie {_METADATA_FILE} missing at {meta_path}. "
                f"Populate /models via the ml-worker downloader."
            )

        with open(meta_path) as f:
            metadata = json.load(f)

        # Per Camie v2 onnx_inference.py: idx_to_tag is keyed by str(idx);
        # tag_to_category maps tag_name -> category. Project to two parallel
        # lists indexed by output position for O(1) lookup in the hot path.
        ds = metadata["dataset_info"]
        idx_to_tag = ds["tag_mapping"]["idx_to_tag"]
        tag_to_category = ds["tag_mapping"]["tag_to_category"]
        total = ds["total_tags"]
        names: list[str] = []
        cats: list[str] = []
        for i in range(total):
            name = idx_to_tag.get(str(i), f"unknown-{i}")
            names.append(name)
            cats.append(tag_to_category.get(name, "general"))

        # Input size from metadata; fall back to 512 (the v2 default).
        self._input_size = int(
            metadata.get("model_info", {}).get("img_size", 512)
        )

        # Lazy import — kept after the file-existence checks so the
        # missing-model RuntimeError still fires first in environments
        # without onnxruntime (CI / lean web image).
        import onnxruntime as ort

        # Cap the intra-op thread pool. ONNX Runtime otherwise sizes it to ALL
        # host cores, so on a shared node each ml-worker replica would grab every
        # core and oversubscribe (and starve the co-located DB/web). Bounding it
        # makes each replica a predictable core consumer — run N replicas where
        # N × _INTRA_OP_THREADS stays within the cores you allot to ML.
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = _INTRA_OP_THREADS
        session = ort.InferenceSession(
            str(model_path), sess_options=opts, providers=["CPUExecutionProvider"],
        )
        self._input_name = session.get_inputs()[0].name
        # Assign sentinels last so a partial load isn't observable.
        self._tag_names = names
        self._tag_categories = cats
        self._session = session

    def _preprocess(self, image_path: Path) -> np.ndarray:
        img = Image.open(image_path)
        # Composite RGBA onto neutral so transparency doesn't bias the model.
        if img.mode == "RGBA":
            bg = Image.new("RGBA", img.size, (255, 255, 255, 255))
            bg.paste(img, mask=img.split()[3])
            img = bg.convert("RGB")
        elif img.mode != "RGB":
            img = img.convert("RGB")

        # Pad to square with ImageNet-mean color, then bicubic resize.
        w, h = img.size
        side = max(w, h)
        square = Image.new("RGB", (side, side), _PAD_COLOR)
        square.paste(img, ((side - w) // 2, (side - h) // 2))
        square = square.resize(
            (self._input_size, self._input_size), Image.BICUBIC
        )

        arr = np.array(square, dtype=np.float32) / 255.0   # HWC, [0,1]
        arr = (arr - _IMAGENET_MEAN) / _IMAGENET_STD        # ImageNet normalize
        arr = arr.transpose(2, 0, 1)                        # HWC -> CHW
        return arr[np.newaxis, :, :, :]                     # NCHW

    def infer(
        self, image_path: Path, *, store_floor: float = DEFAULT_STORE_FLOOR,
    ) -> dict[str, TagPrediction]:
        """Run Camie v2 on one image. Returns {name: TagPrediction} with
        confidence >= store_floor (across all categories — the suggestion
        service does category filtering later). store_floor is the DB-backed
        ml_settings.tagger_store_floor, passed in by the ml task.

        v2 emits multiple outputs; we use the refined predictions
        (output[1] per onnx_inference.py). Sigmoid is applied to raw
        logits to produce [0,1] confidence scores.
        """
        self.load()
        x = self._preprocess(image_path)
        outputs = self._session.run(None, {self._input_name: x})
        # Refined predictions if present (v2 emits initial + refined),
        # fall back to initial for single-output forks.
        logits = outputs[1] if len(outputs) > 1 else outputs[0]
        # Squeeze batch dim, apply sigmoid.
        probs = 1.0 / (1.0 + np.exp(-logits[0]))
        results: dict[str, TagPrediction] = {}
        names = self._tag_names
        cats = self._tag_categories
        for idx, score in enumerate(probs):
            conf = float(score)
            if conf < store_floor:
                continue
            if idx >= len(names):
                # Output longer than metadata declared — shouldn't happen but
                # don't crash the import pipeline if v2 metadata desynchronizes.
                continue
            results[names[idx]] = TagPrediction(
                name=names[idx], category=cats[idx], confidence=conf
            )
        return results


_default_tagger: Tagger | None = None


def get_tagger() -> Tagger:
    """Process-level singleton so the ONNX session loads once per worker."""
    global _default_tagger
    if _default_tagger is None:
        _default_tagger = Tagger()
    return _default_tagger
