"""Custom inference handler for a Dedicated HF Inference Endpoint serving
google/siglip2-so400m-patch14-384 as a pure image embedder.

Returns get_image_features() -> a single pooled, L2-normalized 1152-float vector,
which is exactly what Atlasio's monument_embeddings (vector(1152)) and the
recognize pipeline expect.

Accepts either:
  - raw image bytes with Content-Type image/jpeg|png (HF toolkit passes a PIL image), or
  - JSON {"inputs": "<base64 image>"} or {"inputs": "<http image url>"}.
"""

import base64
import io
from typing import Any, Dict, List, Union

import torch
from PIL import Image
from transformers import AutoModel, AutoProcessor


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.processor = AutoProcessor.from_pretrained(path)
        self.model = AutoModel.from_pretrained(path).to(self.device).eval()

    def _to_image(self, inputs: Any) -> Image.Image:
        if isinstance(inputs, Image.Image):
            return inputs.convert("RGB")
        if isinstance(inputs, (bytes, bytearray)):
            return Image.open(io.BytesIO(bytes(inputs))).convert("RGB")
        if isinstance(inputs, str):
            if inputs.startswith("http://") or inputs.startswith("https://"):
                import urllib.request
                with urllib.request.urlopen(inputs) as r:
                    return Image.open(io.BytesIO(r.read())).convert("RGB")
            # assume base64
            return Image.open(io.BytesIO(base64.b64decode(inputs))).convert("RGB")
        raise ValueError(f"Unsupported input type: {type(inputs)}")

    @torch.no_grad()
    def __call__(self, data: Dict[str, Any]) -> Union[List[float], Dict[str, str]]:
        inputs = data.get("inputs", data)
        image = self._to_image(inputs)
        pixel = self.processor(images=image, return_tensors="pt").to(self.device)
        feats = self.model.get_image_features(**pixel)
        feats = torch.nn.functional.normalize(feats, p=2, dim=-1)
        return feats[0].detach().cpu().tolist()  # length 1152