"""Image inputs.

Qwen3.5 reads images natively: its processor cuts each image into 16-px patches, the vision
encoder merges 2×2 patches into one token, and the language model sees those tokens between
``<|vision_start|>`` and ``<|vision_end|>`` at ``<|image_pad|>`` placeholders, with 3-D (M-RoPE)
positions. ``ImageEncoder`` turns one decision image into the placeholder ids for the layout and
the pixel tensors for ``collate``; ``DecisionModel.hidden`` scatters the encoder's features into
the placeholders.

Glance 0.1 was trained and evaluated on images at most 640 px on the long side;
``encode_image`` prepares an image the same way (downscaled, JPEG at quality 88, as a ``data:`` URL).
"""

from __future__ import annotations

import base64
import io
from pathlib import Path
from typing import Any

from PIL import Image

MAX_SIDE = 640
JPEG_QUALITY = 88


def load_image(entry: dict[str, Any]) -> Image.Image:
    """The image an entry refers to (``image``, ``path`` or ``data_url``), as RGB."""
    if "image" in entry:
        return entry["image"].convert("RGB")
    if "path" in entry:
        with Image.open(Path(entry["path"])) as im:
            return im.convert("RGB")
    if "data_url" in entry:
        data = base64.b64decode(entry["data_url"].split(",", 1)[1])
        with Image.open(io.BytesIO(data)) as im:
            return im.convert("RGB")
    raise ValueError("image entry needs image, path or data_url")


def to_data_url(img: Image.Image, max_side: int | None = MAX_SIDE, quality: int = JPEG_QUALITY) -> str:
    img = img.convert("RGB")
    if max_side and max(img.size) > max_side:
        scale = max_side / max(img.size)
        img = img.resize((max(1, round(img.width * scale)), max(1, round(img.height * scale))), Image.BILINEAR)
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=quality)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode("ascii")


def encode_image(image: "str | Path | Image.Image", max_side: int | None = MAX_SIDE, quality: int = JPEG_QUALITY) -> str:
    """A file path, PIL image or ``data:`` URL as the ``data:`` URL a request carries."""
    if isinstance(image, Image.Image):
        return to_data_url(image, max_side, quality)
    if isinstance(image, str) and image.startswith("data:"):
        return image
    with Image.open(Path(image)) as im:
        return to_data_url(im, max_side, quality)


class ImageEncoder:
    def __init__(self, image_processor: Any, *, image_token_id: int, start_id: int | None = None,
                 end_id: int | None = None, merge_size: int = 2):
        self.image_processor = image_processor
        self.image_token_id = image_token_id
        self.start_id, self.end_id = start_id, end_id
        self.merge = merge_size

    def __call__(self, entry: dict[str, Any]) -> tuple[list[int], dict[str, Any]]:
        """Placeholder ids for one image, and its ``pixel_values`` / ``image_grid_thw``."""
        out = self.image_processor(images=[load_image(entry)], return_tensors="pt")
        grid = out["image_grid_thw"]
        n = int(grid.prod()) // self.merge ** 2
        ids = [self.image_token_id] * n
        if self.start_id is not None:
            ids = [self.start_id, *ids]
        if self.end_id is not None:
            ids = [*ids, self.end_id]
        return ids, {"pixel_values": out["pixel_values"], "image_grid_thw": grid}