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

        self._torch = torch
        self._processor = AutoProcessor.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