"""SigLIP SO400M image-embedding wrapper (PyTorch CPU)."""
from __future__ import annotations
import os

import numpy as np
from PIL import Image

# Defer torch/transformers imports to lazily-loaded functions to allow
# importing MODEL_VERSION in non-ML-worker contexts (e.g., web container
# centroid-recompute enqueue logic).
_torch = None
_AutoModel = None
_AutoProcessor = None

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')
# Model files live flat under this directory (written by scripts/download_models.py via
# snapshot_download(local_dir=...)). We point from_pretrained at the local path directly
# so transformers bypasses its HF cache layout and doesn't need network access at load time.
_LOCAL_DIR = os.path.join(os.environ.get('ML_MODEL_DIR', '/models'), 'siglip')

_model = None
_processor = None


# NOT thread-safe. Must run in the ml-worker container with --concurrency=1.
def _load() -> None:
    global _model, _processor, _torch, _AutoModel, _AutoProcessor
    if _model is not None:
        return
    # Lazy import torch/transformers so this module can be imported in contexts
    # where they're not available (e.g., web container for centroid-recompute enqueue).
    import torch
    from transformers import AutoModel, AutoProcessor
    _torch = torch
    _AutoModel = AutoModel
    _AutoProcessor = AutoProcessor

    _processor = _AutoProcessor.from_pretrained(_LOCAL_DIR)
    _model = _AutoModel.from_pretrained(_LOCAL_DIR)
    _model.eval()


def infer(image_path: str) -> np.ndarray:
    """Return a 1152-dim float32 numpy embedding for the image.

    SigLIP uses a MAP-pooled vision head — the pooled output is the retrieval-ready
    embedding the model was trained to produce. `get_image_features` on transformers
    >= 4.45 returns a BaseModelOutputWithPooling, so pull `.pooler_output` explicitly
    rather than relying on the first-field fallback from indexing.
    """
    _load()
    img = Image.open(image_path).convert('RGB')
    with _torch.no_grad():
        inputs = _processor(images=img, return_tensors='pt')
        out = _model.get_image_features(**inputs)
    pooled = out.pooler_output if hasattr(out, 'pooler_output') else out  # (1, 1152)
    return pooled[0].numpy().astype(np.float32)