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

Runs on CPU by default; moves to CUDA when requested (FC_ML_DEVICE) and a GPU is
available (#872), else stays on CPU. fp32 is kept on GPU too so GPU-computed
embeddings stay in the same numeric space as the existing CPU ones (cosine
comparisons). 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
        self._device = "cpu"

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

        from .device import gpu_requested, torch_mem_fraction

        self._torch = torch
        # GPU (#872) when requested AND a CUDA device is present; else CPU. Cap
        # torch's share of the 8GB P4 (the ONNX tagger shares the card).
        if gpu_requested() and torch.cuda.is_available():
            self._device = "cuda"
            try:
                torch.cuda.set_per_process_memory_fraction(torch_mem_fraction())
            except Exception:  # noqa: BLE001 — best-effort cap; never block load
                pass
        # 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()
        if self._device == "cuda":
            self._model = self._model.to("cuda")

    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")
            if self._device == "cuda":
                inputs = {k: v.to("cuda") for k, v in inputs.items()}
            out = self._model.get_image_features(**inputs)
        pooled = out.pooler_output if hasattr(out, "pooler_output") else out
        # .detach().cpu() so a CUDA tensor converts to numpy (no-op on CPU).
        return pooled[0].detach().cpu().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