# Custom handler for Hugging Face Inference Endpoints.
# Serves imageomics/bioclip-2 (open_clip ViT-L/14, BioCLIP 2, 768-dim embeddings).
#
# API contract (POST JSON to the endpoint):
#   {"inputs": {
#       "image": "<base64-encoded image bytes>",   # or "https://...jpg" URL
#       "candidate_labels": ["Canis familiaris (dog)", ...],   # required for classify
#       "mode": "classify"   # default: zero-shot classification over candidate_labels
#   }}
#   -> [{"label": "...", "score": 0.93}, ...]  sorted by score desc
#
#   {"inputs": {"image": "<base64>", "mode": "embed"}}
#   -> {"embedding": [768 floats]}  (L2-normalized image embedding)
#
#   {"inputs": {"text": "...", "mode": "embed_text"}}
#   -> {"embedding": [768 floats]}  (L2-normalized text embedding)

from typing import Any, Dict, List, Union

import base64
import io
import urllib.request

import open_clip
import torch
import torch.nn.functional as F
from PIL import Image


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # The endpoint mounts the repo contents at `path`; open_clip's
        # "local-dir:" loader reads open_clip_config.json + weights + tokenizer
        # from there. Fall back to the Hub if the mounted path is empty.
        model_id = f"local-dir:{path}" if path else "hf-hub:imageomics/bioclip-2"

        load_kwargs = {}
        if self.device == "cuda":
            load_kwargs["precision"] = "fp16"  # ~0.9GB VRAM instead of ~1.7GB
        try:
            self.model, _, self.preprocess = open_clip.create_model_and_transforms(
                model_id, **load_kwargs
            )
        except TypeError:
            # open_clip versions without the `precision` factory kwarg
            self.model, _, self.preprocess = open_clip.create_model_and_transforms(model_id)

        self.model = self.model.to(self.device).eval()
        self.tokenizer = open_clip.get_tokenizer(model_id)
        print(f"[bioclip-2 handler] loaded {model_id} on {self.device}")

    def _load_image(self, image_field: str) -> Image.Image:
        if image_field.startswith("http://") or image_field.startswith("https://"):
            with urllib.request.urlopen(image_field, timeout=15) as r:
                raw = r.read()
        else:
            if "," in image_field[:64] and image_field[:5].lower() == "data:":
                image_field = image_field.split(",", 1)[1]
            raw = base64.b64decode(image_field)
        return Image.open(io.BytesIO(raw)).convert("RGB")

    def _autocast(self):
        # fp16 weights on CUDA: align input/weight dtypes via autocast —
        # casting the image tensor manually breaks the first conv in open_clip.
        return torch.autocast(
            device_type=self.device, dtype=torch.float16, enabled=self.device == "cuda"
        )

    def _embed_image(self, image_field: str) -> List[float]:
        image = self._load_image(image_field)
        img = self.preprocess(image).unsqueeze(0).to(self.device)
        with self._autocast():
            return F.normalize(self.model.encode_image(img), dim=-1)[0].float().tolist()

    def _embed_text(self, text: str) -> List[float]:
        with self._autocast():
            tok = self.tokenizer([text]).to(self.device)
            return F.normalize(self.model.encode_text(tok), dim=-1)[0].float().tolist()

    def _classify(self, image_field: str, labels: List[str]) -> List[Dict[str, Any]]:
        image = self._load_image(image_field)
        img = self.preprocess(image).unsqueeze(0).to(self.device)
        with self._autocast():
            img_emb = F.normalize(self.model.encode_image(img), dim=-1)
            tok = self.tokenizer(labels).to(self.device)
            txt_emb = F.normalize(self.model.encode_text(tok), dim=-1).float()
            logits = (self.model.logit_scale.exp() * (img_emb @ txt_emb.T)).float()
        # Standard CLIP zero-shot scoring: logit_scale * (img . txt^T), softmax
        # over candidates — the recipe from the BioCLIP model card.
        probs = F.softmax(logits.squeeze(0), dim=0)
        ranked = sorted(zip(labels, probs.tolist()), key=lambda x: -x[1])
        return [{"label": label, "score": float(score)} for label, score in ranked]

    @torch.no_grad()
    def __call__(self, data: Union[Dict[str, Any], List[Dict[str, Any]]]) -> Any:
        # Some handler generations wrap payloads in a list; accept both.
        if isinstance(data, list):
            data = data[0]

        inputs = data.pop("inputs", data)
        mode = inputs.pop("mode", "classify")

        if mode == "embed_text":
            return {"embedding": self._embed_text(inputs["text"])}
        if mode == "embed":
            return {"embedding": self._embed_image(inputs["image"])}
        return self._classify(inputs["image"], inputs["candidate_labels"])