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

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