"""Custom Inference Handler for Hugging Face Inference Endpoints.

This checkpoint is a Qwen3.5 vision-language, tool-calling SFT model
(``Qwen3_5ForConditionalGeneration``). The handler wraps the model in the
``EndpointHandler`` interface expected by the HF Inference Toolkit:

    https://huggingface.co/docs/inference-endpoints/main/en/engines/toolkit

Request payload (``data``) accepts the following keys:

    inputs      (required)  Either:
                              * a plain string prompt, or
                              * a list of OpenAI/Qwen-style chat messages, e.g.
                                [{"role": "user", "content": "Hello"}]
                                Message ``content`` may be a string or a list of
                                {"type": "text"|"image"|"video", ...} parts. Image
                                parts accept "image" / "image_url" as a URL,
                                data URI, or base64 string.
    tools       (optional)  List of tool/function JSON schemas. When provided they
                            are injected via the model's chat template so the model
                            can emit <tool_call> blocks.
    parameters  (optional)  Dict of generation kwargs (max_new_tokens,
                            temperature, top_p, top_k, do_sample, ...).

Response: a list with a single dict ``[{"generated_text": "..."}]``.
"""

from __future__ import annotations

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

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


# Defaults applied when the caller does not override them via ``parameters``.
DEFAULT_GENERATION_KWARGS: Dict[str, Any] = {
    "max_new_tokens": 1024,
    "do_sample": False,
    "temperature": 0.7,
    "top_p": 0.9,
}


class EndpointHandler:
    def __init__(self, path: str = "") -> None:
        # ``path`` is the directory with the model weights provided by the
        # Inference Endpoint runtime (the repository root).
        self.processor = AutoProcessor.from_pretrained(path, trust_remote_code=True)

        dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
        self.model = AutoModelForImageTextToText.from_pretrained(
            path,
            dtype=dtype,
            device_map="auto" if torch.cuda.is_available() else None,
            trust_remote_code=True,
        )
        self.model.eval()

        self.tokenizer = getattr(self.processor, "tokenizer", self.processor)

    # ------------------------------------------------------------------ #
    # Helpers
    # ------------------------------------------------------------------ #
    @staticmethod
    def _load_image(ref: str) -> Image.Image:
        """Load a PIL image from a URL, data URI, or raw base64 string."""
        if ref.startswith("http://") or ref.startswith("https://"):
            import requests

            resp = requests.get(ref, timeout=30)
            resp.raise_for_status()
            return Image.open(io.BytesIO(resp.content)).convert("RGB")

        if ref.startswith("data:"):
            # data:image/png;base64,XXXX
            ref = ref.split(",", 1)[1]

        return Image.open(io.BytesIO(base64.b64decode(ref))).convert("RGB")

    def _collect_images(self, messages: List[Dict[str, Any]]) -> List[Image.Image]:
        images: List[Image.Image] = []
        for message in messages:
            content = message.get("content")
            if not isinstance(content, list):
                continue
            for item in content:
                if not isinstance(item, dict):
                    continue
                if item.get("type") == "image" or "image" in item or "image_url" in item:
                    ref = item.get("image") or item.get("image_url")
                    if isinstance(ref, dict):
                        ref = ref.get("url")
                    if isinstance(ref, str):
                        images.append(self._load_image(ref))
        return images

    def _build_messages(
        self, inputs: Union[str, List[Dict[str, Any]]]
    ) -> List[Dict[str, Any]]:
        if isinstance(inputs, str):
            return [{"role": "user", "content": inputs}]
        return inputs

    # ------------------------------------------------------------------ #
    # Inference
    # ------------------------------------------------------------------ #
    @torch.inference_mode()
    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        inputs = data.pop("inputs", data)
        tools: Optional[List[Dict[str, Any]]] = data.pop("tools", None)
        parameters: Dict[str, Any] = data.pop("parameters", {}) or {}

        messages = self._build_messages(inputs)
        images = self._collect_images(messages)

        prompt = self.processor.apply_chat_template(
            messages,
            tools=tools,
            tokenize=False,
            add_generation_prompt=True,
        )

        processor_kwargs: Dict[str, Any] = {"text": [prompt], "return_tensors": "pt"}
        if images:
            processor_kwargs["images"] = images

        model_inputs = self.processor(**processor_kwargs)
        model_inputs = {k: v.to(self.model.device) for k, v in model_inputs.items()}

        generation_kwargs = {**DEFAULT_GENERATION_KWARGS, **parameters}
        # When sampling is disabled, drop sampling-only knobs to avoid warnings.
        if not generation_kwargs.get("do_sample", False):
            for key in ("temperature", "top_p", "top_k"):
                generation_kwargs.pop(key, None)

        generated_ids = self.model.generate(**model_inputs, **generation_kwargs)

        # Strip the prompt tokens so only the completion is decoded.
        input_len = model_inputs["input_ids"].shape[1]
        new_tokens = generated_ids[:, input_len:]

        text = self.processor.batch_decode(
            new_tokens,
            skip_special_tokens=True,
            clean_up_tokenization_spaces=False,
        )[0]

        return [{"generated_text": text}]