"""SigLIP SO400M image-embedding wrapper (PyTorch CPU).

torch/transformers are imported lazily inside load() so this module can be
imported in the web container (which never runs inference) without paying the
torch import cost.
"""

import os
from pathlib import Path

import numpy as np
from PIL import Image, ImageFile

from ..worker_lanes import LANES_BY_NAME

ImageFile.LOAD_TRUNCATED_IMAGES = True

# Cap torch's intra-op threads so each ml-worker replica is a bounded core
# consumer on a shared node (torch otherwise uses all cores).
#
# Read from the lane rather than restated here. This was a literal 4 beside a
# comment reading "keep N_replicas x this within the cores allotted to ML" —
# a constraint written where nothing could act on it, and nothing did: the ML
# ceiling came from memory alone, offered the operator ~49 slots on a
# large-memory host, and the lane spent 2026-09-23 with ~200 torch threads on
# it. `derived_ceiling` now divides the cores by this number, which only means
# anything while the two are the same number.
_INTRA_OP_THREADS = LANES_BY_NAME["ml"].threads_per_slot

DEFAULT_MODEL_NAME = os.environ.get(
    "SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
)
# Back-compat alias (api/gpu imported this name as the fallback embedder id).
MODEL_NAME = DEFAULT_MODEL_NAME
MODEL_VERSION = os.environ.get(
    "SIGLIP_MODEL_VERSION", "siglip-so400m-patch14-384"
)
EMBED_DIM = 1152
_LOCAL_DIR = Path(os.environ.get("ML_MODEL_DIR", "/models")) / "siglip"


class Embedder:
    """Loads whatever SigLIP-family model it's given by HF NAME. For the default
    model it prefers the pre-downloaded local dir (no re-download on existing
    deploys); any other name resolves as an HF repo id (downloaded + cached on
    first use), so an operator model swap (#1190) just works server-side."""

    def __init__(self, model_name: str | None = None, model_dir: Path | None = None):
        self.model_name = model_name or DEFAULT_MODEL_NAME
        self._explicit_dir = model_dir
        self._model = None
        self._processor = None
        self._torch = None

    def _source(self) -> str:
        if self._explicit_dir is not None:
            return str(self._explicit_dir)
        if self.model_name == DEFAULT_MODEL_NAME and _LOCAL_DIR.exists():
            return str(_LOCAL_DIR)
        return self.model_name

    def load(self) -> None:
        if self._model is not None:
            return
        import torch
        from transformers import AutoImageProcessor, AutoModel

        self._torch = torch
        # Bound torch's CPU thread pool (see _INTRA_OP_THREADS) so each replica
        # stays a predictable core consumer on a shared node.
        torch.set_num_threads(_INTRA_OP_THREADS)
        # IMAGE inference only — AutoImageProcessor loads just the image side
        # (preprocessor_config.json), skipping the SigLIP tokenizer + its
        # sentencepiece dep (operator hit that ImportError 2026-05-25). Works
        # for any SigLIP-family model, keeping the embedder model-agnostic.
        src = self._source()
        self._processor = AutoImageProcessor.from_pretrained(src)
        self._model = AutoModel.from_pretrained(src)
        self._model.eval()

    def infer(self, image_path: Path) -> np.ndarray:
        """Return a 1152-dim float32 embedding (SigLIP MAP-pooled output)."""
        self.load()
        img = Image.open(image_path).convert("RGB")
        with self._torch.no_grad():
            inputs = self._processor(images=img, return_tensors="pt")
            out = self._model.get_image_features(**inputs)
        pooled = out.pooler_output if hasattr(out, "pooler_output") else out
        return pooled[0].numpy().astype(np.float32)


_default_embedder: Embedder | None = None


def get_embedder(model_name: str | None = None) -> Embedder:
    """Cached embedder for `model_name` (default if None). Rebuilds the singleton
    when the requested name changes, so an operator model swap takes effect
    without restarting the worker."""
    global _default_embedder
    name = model_name or DEFAULT_MODEL_NAME
    if _default_embedder is None or _default_embedder.model_name != name:
        _default_embedder = Embedder(model_name=name)
    return _default_embedder