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

Direct port of ImageRepo's siglip.py. 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

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