"""
Custom Hugging Face Inference Endpoint handler for Qwen3-VL-Embedding-8B.

WHERE THIS FILE GOES:
Hugging Face Inference Endpoints look for a `handler.py` file living in the
ROOT of the MODEL REPO you deploy (not in your own project repo). So:

  1. Duplicate Qwen/Qwen3-VL-Embedding-8B into your own namespace on the Hub
     (huggingface.co -> the model page -> "..." menu -> "Duplicate this model"),
     e.g. your-username/qwen3-vl-embedding-endpoint. This copies the weights
     without you having to re-upload ~16GB yourself.
  2. Add this file to that new repo, named exactly `handler.py`, in the repo root.
  3. Add the accompanying `requirements.txt` (see endpoint_requirements.txt)
     to that same repo root.
  4. Deploy an Inference Endpoint from that repo. Because it contains a
     handler.py, Endpoints will use it automatically instead of a default
     pipeline.

WHAT IT DOES:
Loads the model once when the endpoint starts, then on every request:
builds the same system-instruction + text/image conversation your old
local code used, runs a forward pass, takes last-token pooling, L2-normalizes,
and returns the embedding as JSON.

REQUEST FORMAT (what your client should POST):
  {"inputs": {"text": "some product text", "image_base64": "<optional b64>"}}

RESPONSE FORMAT:
  {"embedding": [0.01, -0.02, ...], "dimension": 4096}
"""
from __future__ import annotations

import base64
from io import BytesIO
from typing import Any, Dict, List, Optional

import torch
import torch.nn.functional as F
from PIL import Image
from transformers import AutoModel, AutoProcessor

INSTRUCTION = (
    "Represent this retail product for multimodal fashion, beauty, and home catalog retrieval. "
    "Preserve product identity, category, materials, visible design details, structure, color nuance, and style-relevant attributes."
)

MAX_IMAGE_SIDE = 768


def _format_as_conversation(text: str, has_image: bool) -> List[Dict[str, Any]]:
    content: List[Dict[str, Any]] = []
    if has_image:
        content.append({"type": "image"})
    if text:
        content.append({"type": "text", "text": text})
    if not content:
        content.append({"type": "text", "text": ""})
    return [
        {"role": "system", "content": [{"type": "text", "text": INSTRUCTION}]},
        {"role": "user", "content": content},
    ]


class EndpointHandler:
    def __init__(self, path: str = ""):
        # `path` is filled in by the Endpoints runtime with the local
        # directory the repo (including model weights) was downloaded into.
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.dtype = torch.bfloat16 if self.device == "cuda" else torch.float32

        self.processor = AutoProcessor.from_pretrained(
            path, trust_remote_code=True, local_files_only=True
        )
        self.model = AutoModel.from_pretrained(
            path, trust_remote_code=True, local_files_only=True, torch_dtype=self.dtype
        )
        self.model = self.model.to(self.device).eval()

    def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
        payload = data.get("inputs", data) or {}
        text = payload.get("text") or ""
        image_b64 = payload.get("image_base64")

        image = self._decode_image(image_b64) if image_b64 else None
        _validate(text, image)

        vector = self._compute(text, image)
        return {"embedding": vector, "dimension": len(vector)}

    def _compute(self, text: str, image: Optional[Image.Image]) -> List[float]:
        images = [image] if image is not None else None
        conversation = _format_as_conversation(text, image is not None)
        prompt_text = self.processor.apply_chat_template(
            conversation, tokenize=False, add_generation_prompt=True
        )
        inputs = self.processor(text=[prompt_text], images=images, padding=True, return_tensors="pt")
        inputs = {key: value.to(self.device) for key, value in inputs.items()}

        with torch.inference_mode():
            outputs = self.model(**inputs)
            hidden = outputs.last_hidden_state
            attention_mask = inputs["attention_mask"]
            last_token_index = attention_mask.sum(dim=1) - 1
            embedding = hidden[0, last_token_index[0]]
            embedding = F.normalize(embedding, p=2, dim=0)

        return embedding.detach().cpu().float().tolist()

    @staticmethod
    def _decode_image(image_b64: str) -> Image.Image:
        try:
            image = Image.open(BytesIO(base64.b64decode(image_b64))).convert("RGB")
        except Exception as exc:
            raise ValueError(f"Invalid image_base64: {exc}") from exc
        image.thumbnail((MAX_IMAGE_SIDE, MAX_IMAGE_SIDE), Image.Resampling.BICUBIC)
        return image


def _validate(text: str, image: Optional[Image.Image]) -> None:
    if not text and image is None:
        raise ValueError("Either 'text' or 'image_base64' must be provided")