"""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

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). Keep
# N_replicas × this within the cores allotted to ML to avoid oversubscription.
_INTRA_OP_THREADS = 4

MODEL_NAME = os.environ.get(
    "SIGLIP_MODEL_NAME", "google/siglip-so400m-patch14-384"
)
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:
    def __init__(self, model_dir: Path | None = None):
        self._model_dir = model_dir or _LOCAL_DIR
        self._model = None
        self._processor = None
        self._torch = None

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

        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)
        # FC's embedder only does IMAGE inference — never text. AutoProcessor
        # loads the full processor including SiglipTokenizer, which requires
        # the sentencepiece library at import time even if we never call it.
        # SiglipImageProcessor loads ONLY preprocessor_config.json (image
        # side) and skips the tokenizer config entirely. Operator hit the
        # ImportError 2026-05-25 once the ml-worker started actually running
        # tag_and_embed; switching to the image-only loader avoids the
        # tokenizer dep without adding ~30 MB of unused C++ build to the
        # lean ml-worker image.
        self._processor = SiglipImageProcessor.from_pretrained(
            str(self._model_dir)
        )
        self._model = AutoModel.from_pretrained(str(self._model_dir))
        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() -> Embedder:
    global _default_embedder
    if _default_embedder is None:
        _default_embedder = Embedder()
    return _default_embedder
