"""OpenAI-multimodal processor wiring for DeepSeek-V4-Flash-Vision-Exp.

dots3_note-style processor trio (workflow of the b12x vLLM fork):

- a small HF-like processor that tokenizes prompt text (placeholders intact)
  and precomputes ViT patches / grid / sentinel layout per image,
- a dummy builder for MM profiling,
- a `BaseMultiModalProcessor` whose prompt update expands every
  `<｜deepseek_image|>` token (vocab id 129264) into the per-image sentinel
  block built by `image_processor.build_image_block`, with the whole block
  flagged as mm embeddings.

`model.embed_multimodal` turns the tensors into per-image embedding runs and
`model.embed_input_ids` scatters them over the flagged range on prefill.
"""

from __future__ import annotations

import math
from collections.abc import Mapping, Sequence
from functools import cached_property
from typing import Any, cast

import numpy as np
import torch
from PIL import Image
from transformers import BatchFeature

from vllm.config.multimodal import BaseDummyOptions, ImageDummyOptions
from vllm.inputs import MultiModalDataDict
from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
from vllm.multimodal.parse import MultiModalDataItems
from vllm.multimodal.processing import (
    BaseDummyInputsBuilder,
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    PromptReplacement,
    PromptUpdateDetails,
)

from .encoding_dsv4 import IMAGE_PLACEHOLDER
from .image_processor import (
    COMPRESS_PAD_TO,
    build_image_block,
    grid_tokens,
    safe_resize,
    solve_resize_ratio,
)

_PROTO_PATCHES = 60  # prototype patch-grid side for the dummy-size sweep


def _vision_kv(hf: Any) -> dict[str, int | float]:
    return {
        "vision_patch_size": int(getattr(hf, "vision_patch_size", 14)),
        "vision_downsample_ratio": int(getattr(hf, "vision_downsample_ratio", 3)),
        "vision_max_n_token": int(getattr(hf, "vision_max_n_token", 384)),
        "vision_min_pixels": int(getattr(hf, "vision_min_pixels", 147456)),
        "vision_max_wh_ratio": int(getattr(hf, "vision_max_wh_ratio", 8)),
    }


class DSV4VisionImagePreparer:
    """Resize/pad/patch a PIL image per the day-0 `load_image` math."""

    def __init__(self, hf: Any) -> None:
        for key, value in _vision_kv(hf).items():
            setattr(self, key, value)
        self.image_token_id: int | None = None

    @property
    def p(self) -> int:
        return self.vision_patch_size

    def prepare(self, image: Image.Image) -> dict[str, torch.Tensor]:
        image = image.convert("RGB")
        width, height = image.size
        if self.vision_max_wh_ratio and width > height * self.vision_max_wh_ratio:
            width = height * self.vision_max_wh_ratio
        if 0 < width * height < self.vision_min_pixels:
            scale = (self.vision_min_pixels / (width * height)) ** 0.5
            width = int(width * scale)
            height = int(height * scale)
        best_width = math.ceil(width / self.p) * self.p
        best_height = math.ceil(height / self.p) * self.p
        n_llm_h, n_llm_w, best_height, best_width = safe_resize(
            height,
            width,
            best_height,
            best_width,
            self.p,
            self.vision_downsample_ratio,
            self.vision_max_n_token,
        )
        n_vit_h, n_vit_w = best_height // self.p, best_width // self.p
        if self.vision_max_wh_ratio and image.width >= self.vision_max_wh_ratio * image.height:
            resized = image.resize((best_width, best_height))
        else:
            resized = _pad_image(image, best_width, best_height)
        x = torch.from_numpy(np.asarray(resized, dtype="float32")).permute(2, 0, 1) / 255
        x = ((x - 0.5) / 0.5).to(torch.bfloat16)
        patches = (
            x.reshape(3, n_vit_h, self.p, n_vit_w, self.p)
            .permute(1, 3, 0, 2, 4)
            .reshape(n_vit_h * n_vit_w, 3, self.p, self.p)
        )
        return dict(
            patches=patches,
            n_vit_h=n_vit_h,
            n_vit_w=n_vit_w,
            n_llm_h=n_llm_h,
            n_llm_w=n_llm_w,
        )

    def max_block_scan(self) -> tuple[int, int, int]:
        """(best_h, best_w) and its block length over candidate aspect ratios."""
        base = _PROTO_PATCHES * self.p * self.vision_downsample_ratio
        ratios = [1 / r for r in (1, 2, 3, 4, 5, 6, 7, 8)]
        ratios += [1 / r for r in ratios[1:]]
        best: tuple[int, int, int] = (0, 0, 0)
        for height_ratio in ratios:
            height = max(self.p * self.vision_downsample_ratio, int(base * height_ratio))
            width = max(self.p * self.vision_downsample_ratio, int(base / height_ratio))
            best_h, best_w = math.ceil(height / self.p) * self.p, math.ceil(width / self.p) * self.p
            n_h, n_w, fb_h, fb_w, _ = solve_resize_ratio(
                height, width, self.p, self.vision_downsample_ratio, self.vision_max_n_token
            )
            best_h, best_w = max(best_h, fb_h), max(best_w, fb_w)
            _, _, num = grid_tokens(fb_h, fb_w, self.p, self.vision_downsample_ratio)
            if num > best[2]:
                best = (fb_h, fb_w, num + COMPRESS_PAD_TO - 1)
        return best


def _pad_image(image: Image.Image, width: int, height: int) -> Image.Image:
    from PIL import ImageOps

    return ImageOps.pad(image, (width, height), color=(127, 127, 127))


class DSV4VisionProcessingInfo(BaseProcessingInfo):
    @cached_property
    def preparer(self) -> DSV4VisionImagePreparer:
        return DSV4VisionImagePreparer(self.get_hf_config())

    @cached_property
    def processor(self) -> "DSV4VisionProcessor":
        return DSV4VisionProcessor(
            self.get_tokenizer(),
            self.preparer,
            vocab_size=int(self.get_hf_config().vocab_size),
        )

    def get_hf_processor(self, **kwargs: object) -> "DSV4VisionProcessor":
        return self.processor

    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
        return {"image": None}

    @cached_property
    def max_image_size(self) -> tuple[int, int]:
        best_h, best_w, _ = self.preparer.max_block_scan()
        return best_w, best_h

    def get_mm_max_tokens_per_item(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
    ) -> Mapping[str, int]:
        if not mm_counts.get("image"):
            return {}
        _, _, num = self.preparer.max_block_scan()
        return {"image": num + 4}


class DSV4VisionProcessor:
    """HF-like processor: tokenize text, precompute image tensors."""

    def __init__(self, tokenizer, preparer: DSV4VisionImagePreparer, vocab_size: int) -> None:
        self.tokenizer = tokenizer
        self.preparer = preparer
        self.vocab_size = vocab_size
        ids = tokenizer.convert_tokens_to_ids(IMAGE_PLACEHOLDER)
        if ids is None or ids == getattr(tokenizer, "unk_token_id", -1):
            raise ValueError(f"Token not found in tokenizer: {IMAGE_PLACEHOLDER}")
        self.image_token_id = int(ids)
        preparer.image_token_id = self.image_token_id

    def _encode(self, text: str) -> list[int]:
        try:
            ids = self.tokenizer.encode(text, add_special_tokens=False)
        except TypeError:
            ids = self.tokenizer.encode(text)
        if ids and isinstance(ids[0], list):
            ids = ids[0]
        return list(cast("list[int]", ids))

    def __call__(self, text: str, images: Sequence[Image.Image] | None = None, **_kw: Any) -> BatchFeature:
        image_list = list(images or [])
        token_ids = self._encode(text)
        found = token_ids.count(self.image_token_id)
        if found < len(image_list):
            extra = " ".join([IMAGE_PLACEHOLDER] * (len(image_list) - found))
            text = f"{text}\n{extra}" if found else text + extra
            token_ids = self._encode(text)
            found = token_ids.count(self.image_token_id)
        if found != len(image_list):
            raise ValueError(
                f"Found {found} image tokens but got {len(image_list)} images"
            )

        patches: list[torch.Tensor] = []
        grids: list[torch.Tensor] = []
        types: list[torch.Tensor] = []
        perms: list[torch.Tensor] = []
        perm_counts: list[int] = []
        running = 0
        image_idx = 0
        for pos, tok in enumerate(token_ids):
            if tok != self.image_token_id:
                running += 1
                continue
            prep = self.preparer.prepare(image_list[image_idx])
            n_vit_h, n_vit_w, n_llm_h, n_llm_w = (
                prep["n_vit_h"],
                prep["n_vit_w"],
                prep["n_llm_h"],
                prep["n_llm_w"],
            )
            blk_types, perm = build_image_block(n_llm_h, n_llm_w, running)
            patches.append(prep["patches"])
            grids.append(
                torch.tensor([n_vit_h, n_vit_w, n_llm_h, n_llm_w], dtype=torch.long)
            )
            types.append(blk_types)
            perms.append(perm)
            perm_counts.append(int(perm.numel()))
            image_idx += 1
            running += int(blk_types.numel())

        data: dict[str, Any] = {"input_ids": [token_ids]}
        if patches:
            data["image_patches"] = torch.cat(patches)
            data["image_grids"] = torch.stack(grids)
            data["image_types"] = torch.cat(types)
            data["image_perm"] = torch.cat(perms)
            data["image_patch_counts"] = torch.tensor(
                [int(pt.shape[0]) for pt in patches], dtype=torch.long
            )
            data["image_type_counts"] = torch.tensor(
                [int(t.numel()) for t in types], dtype=torch.long
            )
            data["image_perm_counts"] = torch.tensor(perm_counts, dtype=torch.long)
        return BatchFeature(data=data)


class DSV4VisionDummyInputsBuilder(BaseDummyInputsBuilder[DSV4VisionProcessingInfo]):
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        return IMAGE_PLACEHOLDER * mm_counts.get("image", 0)

    def get_dummy_mm_data(
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
        mm_options: Mapping[str, BaseDummyOptions],
    ) -> MultiModalDataDict:
        num = mm_counts.get("image", 0)
        if not num:
            return {}
        width, height = self.info.max_image_size
        overrides = cast(ImageDummyOptions | None, mm_options.get("image"))
        if overrides is not None and overrides.width and overrides.height:
            width, height = overrides.width, overrides.height
        image = Image.new("RGB", (width, height), color=(127, 127, 127))
        return {"image": [image] * num}


class DSV4VisionMultiModalProcessor(BaseMultiModalProcessor[DSV4VisionProcessingInfo]):
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
        tok_kwargs: Mapping[str, object],
    ) -> BatchFeature:
        return self.info.ctx.call_hf_processor(
            self.info.get_hf_processor(**mm_kwargs),
            dict(text=prompt, **mm_data),
            dict(**mm_kwargs, **tok_kwargs),
        )

    def _hf_processor_applies_updates(
        self,
        prompt_text: str,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        tokenization_kwargs: Mapping[str, object],
    ) -> bool:
        return False

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        if "image_patches" not in hf_inputs:
            return {}
        patch_counts = hf_inputs["image_patch_counts"]
        type_counts = hf_inputs["image_type_counts"]
        perm_counts = hf_inputs["image_perm_counts"]
        ones = torch.ones(len(patch_counts), dtype=torch.long)
        fields: dict[str, MultiModalFieldConfig] = {
            "image_patches": MultiModalFieldConfig.flat_from_sizes(
                "image", patch_counts
            ),
            "image_perm": MultiModalFieldConfig.flat_from_sizes(
                "image", perm_counts
            ),
            "image_types": MultiModalFieldConfig.flat_from_sizes(
                "image", type_counts
            ),
        }
        for key in (
            "image_grids",
            "image_patch_counts",
            "image_type_counts",
            "image_perm_counts",
        ):
            fields[key] = MultiModalFieldConfig.flat_from_sizes("image", ones)
        return fields

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
        out_mm_kwargs: MultiModalKwargsItems,
    ) -> Sequence:
        token_id = self.info.processor.image_token_id

        def replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
            count = int(out_mm_kwargs["image"][item_idx]["image_type_counts"].data)
            return PromptUpdateDetails.select_token_id(
                [token_id] * count, token_id
            )

        return [
            PromptReplacement(
                modality="image",
                target=[token_id],
                replacement=replacement,
            )
        ]