# Custom inference handler for the palatshq DINOv2 image-embedding endpoint.
#
# This file is the source of truth for the Hugging Face model repo
# (palatshq/dinov2-large). It is NOT imported by the API build; it lives here so
# the deployed code is version-controlled alongside its client
# (apps/api/src/lib/dino-embedder.ts). Push it to the HF repo's `main` branch and
# redeploy the endpoint — see docs/integrations/huggingface.md.
#
# DINOv2 is image-only (no text tower). We return the CLS-token embedding
# (Dinov2 pooler_output), L2-normalized so cosine similarity is a plain dot
# product — matching how the CLIP endpoint returns normalized vectors and what
# the benchmark's cosine matcher expects.
#
# Contract (mirrors dino-embedder.ts). Single image:
#   request:  { "inputs": { "image": "<base64 or data-url>" } }
#   response: { "image_embedding": [float, ...] }
# Batch (one forward over many images, for backfills):
#   request:  { "inputs": { "images": ["<base64>", ...] } }
#   response: { "image_embeddings": [[float, ...], ...] }  # aligned with input order

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

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


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

    @staticmethod
    def _load_image(image_value: str) -> Image.Image:
        if image_value.startswith("data:"):
            image_value = image_value.split(",", 1)[1]
        try:
            raw = base64.b64decode(image_value, validate=True)
        except (binascii.Error, ValueError) as error:
            raise ValueError("`image` must be base64 or a base64 data-url") from error
        return Image.open(io.BytesIO(raw)).convert("RGB")

    def _embed(self, images: List[Image.Image]) -> List[List[float]]:
        processed = self.processor(images=images, return_tensors="pt").to(self.device)

        # fp16 autocast on GPU roughly halves the vision-encoder forward cost
        # with no meaningful hit to embedding quality; CPU stays fp32 (half is
        # slow/unsupported for some ops there).
        if self.device == "cuda":
            with torch.autocast("cuda", dtype=torch.float16):
                outputs = self.model(**processed)
        else:
            outputs = self.model(**processed)

        # pooler_output is the CLS token after the final layernorm — the standard
        # global descriptor for DINOv2 retrieval.
        embeddings = outputs.pooler_output.float()
        embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True)
        return embeddings.cpu().tolist()

    @torch.inference_mode()
    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        inputs = data.get("inputs", data)
        if not isinstance(inputs, dict):
            raise ValueError("Request must contain `inputs`")

        raw_images = inputs.get("images")
        if raw_images is not None:
            if not isinstance(raw_images, list) or len(raw_images) == 0:
                raise ValueError("`inputs.images` must be a non-empty list")
            images = [self._load_image(value) for value in raw_images]
            return {"image_embeddings": self._embed(images)}

        if inputs.get("image"):
            image = self._load_image(inputs["image"])
            return {"image_embedding": self._embed([image])[0]}

        raise ValueError("Request must contain `inputs.image` or `inputs.images`")