"""Official DeepSeek-V4-Flash-Vision-Exp image preprocessor.

Adapted from deepseek-ai/DeepSeek-V4-Flash-Vision-Exp ``inference/image_processor.py``
(MIT). Adds a PIL entry point so vLLM can feed decoded images without going
back through URLs. Video is not part of this checkpoint: GIF is decoded as a
still RGB frame (PIL's first frame).
"""
from __future__ import annotations

import base64
import io
import math
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any, TYPE_CHECKING
from urllib.request import urlopen

if TYPE_CHECKING:
    import torch
    from PIL import Image as PILImage


def _pil():
    """Load Pillow on first image decode. CPU CI does not install it."""
    try:
        from PIL import Image, ImageOps
    except ImportError as exc:
        raise ModuleNotFoundError(
            "Pillow is required to decode Vision-Exp images"
        ) from exc
    return Image, ImageOps


IMAGE_START, IMAGE_PAD, IMAGE, IMAGE_NEW_LINE, IMAGE_END = range(5)
COMPRESS_PAD_TO = 4
IMAGE_PLACEHOLDER = "<｜deepseek_image｜>"


def is_vision_exp_weight_name(name: str) -> bool:
    """Checkpoint keys for the ViT/Aligner/image tokens.

    Anemll ``load_weights`` treats any name containing ``w1`` as a stacked
    ``gate_up_proj`` shard. Aligner's ``w1`` and ViT MLP ``w1`` are full
    Linear layers, so those keys must bypass the stacked mapping.
    """
    return (
        name.startswith("aligner.")
        or name.startswith("vision.")
        or name.startswith("image_")
        or name.startswith("model.aligner.")
        or name.startswith("model.vision.")
        or name.startswith("model.image_")
    )


def is_unregistered_router_bias(name: str, param_names: Any) -> bool:
    """True when a gate bias has no Parameter.

    Hash MoE (layers 0–2) does not register ``e_score_correction_bias``.
    0731 omitted those tensors; Vision-Exp dumps ``ffn.gate.bias`` /
    ``bias_vl`` there anyway. Anemll's mapper rewrites them to
    ``e_score_correction_bias`` / ``_vl``, then ``load_weights`` KeyErrors.
    The DSpark draft loader uses ``endswith(".ffn.gate.bias")``, which
    misses ``bias_vl``, so unmapped ``.ffn.gate.bias_vl`` is also unused
    until that remap is patched.
    """
    if name in param_names:
        return False
    return name.endswith(
        (
            ".ffn.gate.e_score_correction_bias",
            ".ffn.gate.e_score_correction_bias_vl",
            ".ffn.gate.bias_vl",
            ".ffn.gate.bias",
        )
    )


def as_pil(item: Any) -> PILImage.Image:
    """Normalize vLLM image items (PIL, HWC/CHW array, tensor, dict) to RGB PIL."""
    Image, _ImageOps = _pil()
    if isinstance(item, Image.Image):
        return item.convert("RGB")
    if isinstance(item, dict):
        for key in ("image_pil", "pil", "image"):
            value = item.get(key)
            if value is None or isinstance(value, dict):
                continue
            return as_pil(value)
        raise TypeError(f"Unsupported image dict keys: {sorted(item)!r}")
    array = item
    try:
        import torch
    except ImportError:
        torch = None  # type: ignore[assignment]
    if torch is not None and isinstance(array, torch.Tensor):
        array = array.detach().cpu().numpy()
    try:
        import numpy as np
    except ImportError as exc:
        raise TypeError(f"Unsupported image item type: {type(item)!r}") from exc
    if not hasattr(array, "ndim"):
        raise TypeError(f"Unsupported image item type: {type(item)!r}")
    arr = np.asarray(array)
    if arr.ndim == 2:
        arr = np.stack([arr, arr, arr], axis=-1)
    elif arr.ndim == 3 and arr.shape[-1] not in (1, 3, 4) and arr.shape[0] in (1, 3, 4):
        arr = np.transpose(arr, (1, 2, 0))
    if arr.ndim != 3:
        raise TypeError(f"Unsupported image array shape: {arr.shape!r}")
    if arr.shape[-1] == 1:
        arr = np.repeat(arr, 3, axis=-1)
    if np.issubdtype(arr.dtype, np.floating):
        peak = float(arr.max()) if arr.size else 0.0
        if peak <= 1.0:
            arr = arr * 255.0
        arr = np.clip(arr, 0, 255).astype("uint8")
    else:
        arr = np.clip(arr, 0, 255).astype("uint8")
    if arr.shape[-1] == 4:
        return Image.fromarray(arr, mode="RGBA").convert("RGB")
    return Image.fromarray(arr[..., :3], mode="RGB")


@dataclass
class ImageInput:
    start: int
    patches: Any
    n_vit_h: int
    n_vit_w: int
    types: Any
    perm: Any


def vision_args_from_config(config: Any) -> SimpleNamespace:
    return SimpleNamespace(
        vision_patch_size=int(getattr(config, "vision_patch_size", 14)),
        vision_downsample_ratio=int(getattr(config, "vision_downsample_ratio", 3)),
        vision_max_n_token=int(getattr(config, "vision_max_n_token", 384)),
        vision_min_pixels=int(getattr(config, "vision_min_pixels", 147456)),
        vision_max_wh_ratio=getattr(config, "vision_max_wh_ratio", 8),
        vision_n_layers=int(getattr(config, "vision_n_layers", 0)),
        vision_dim=int(getattr(config, "vision_dim", 1024)),
        vision_n_heads=int(getattr(config, "vision_n_heads", 16)),
        vision_inter_dim=int(getattr(config, "vision_inter_dim", 2816)),
        vision_rope_theta=float(getattr(config, "vision_rope_theta", 10000.0)),
        vocab_size=int(getattr(config, "vocab_size", 129280)),
        dim=int(getattr(config, "hidden_size", getattr(config, "dim", 4096))),
        hidden_size=int(getattr(config, "hidden_size", getattr(config, "dim", 4096))),
    )


def grid_tokens(best_height, best_width, patch_size, downsample_ratio):
    """Number of LLM tokens the aligner grid occupies (N-layout, incl. row/align padding)."""
    n_llm_h = math.ceil((best_height // patch_size) / downsample_ratio)
    n_llm_w = math.ceil((best_width // patch_size) / downsample_ratio)
    num_tokens = n_llm_h * (n_llm_w + 1) + 2
    if n_llm_h % 2 == 1:
        num_tokens += n_llm_w + 1
    num_tokens += (n_llm_h + 1) // 2 * (n_llm_w + 1) % 2 * 2
    return n_llm_h, n_llm_w, num_tokens


def solve_resize_ratio(height, width, patch_size, downsample_ratio, max_n_token):
    r = height / width
    max_w_float = math.sqrt((max_n_token - 2) / r + 0.25) - 0.5
    max_h_float = max_w_float * r
    if max_w_float < 1.0:
        max_w = 1
        max_h = (max_n_token - 2) // (max_w + 1)
        if max_h % 2 == 1:
            max_h -= 1
        best_width = max_w * patch_size * downsample_ratio
        best_height = max_h * patch_size * downsample_ratio
    elif max_h_float < 2.0:
        max_h = 2
        max_w = ((max_n_token - 2) // max_h) - 1
        assert max_w > 1
        best_width = max_w * patch_size * downsample_ratio
        best_height = max_h * patch_size * downsample_ratio
    else:
        max_w = math.floor(max_w_float)
        max_h = math.floor(max_h_float)
        if max_h % 2 == 1:
            max_h -= 1
        beta = min(
            max_w * patch_size * downsample_ratio / width,
            max_h * patch_size * downsample_ratio / height,
        )
        best_width = math.floor(width * beta / patch_size) * patch_size
        best_height = math.floor(height * beta / patch_size) * patch_size
    n_llm_h, n_llm_w, num_tokens = grid_tokens(
        best_height, best_width, patch_size, downsample_ratio
    )
    return n_llm_h, n_llm_w, best_height, best_width, num_tokens


def safe_resize(height, width, best_height, best_width, patch_size, downsample_ratio, max_n_token):
    max_n_token -= COMPRESS_PAD_TO - 1
    n_llm_h, n_llm_w, num_tokens = grid_tokens(
        best_height, best_width, patch_size, downsample_ratio
    )
    budget = max_n_token
    while num_tokens > max_n_token:
        n_llm_h, n_llm_w, best_height, best_width, num_tokens = solve_resize_ratio(
            height, width, patch_size, downsample_ratio, budget
        )
        budget -= 1
    return n_llm_h, n_llm_w, best_height, best_width


def load_image_bytes(record) -> bytes:
    """Load image bytes from raw/base64 data, an Anthropic source, URL, or path."""
    data = record.get("data")
    if isinstance(data, bytes):
        return data
    if isinstance(data, str):
        return base64.b64decode(data)

    source = record.get("source")
    if isinstance(source, dict):
        if source.get("data") is not None:
            return base64.b64decode(source["data"])
        if source.get("url"):
            return load_image_bytes({"url": source["url"]})

    url = record.get("url")
    if isinstance(url, str) and url:
        if url.startswith("data:"):
            header, _, payload = url.partition(",")
            if ";base64" not in header:
                raise ValueError(f"Unsupported data URL encoding: {header}")
            return base64.b64decode(payload)
        if url.startswith(("http://", "https://")):
            with urlopen(url, timeout=30) as response:
                return response.read()
        with open(url, "rb") as file:
            return file.read()

    raise ValueError(f"Cannot load image from record: {list(record.keys())}")


def pil_to_patches(image: PILImage.Image, args) -> tuple[Any, int, int, int, int]:
    """Resize/pad one RGB image and return ViT patches plus LLM grid sizes."""
    import torch

    _Image, ImageOps = _pil()

    p = args.vision_patch_size
    image = image.convert("RGB")
    width, height = image.size
    if args.vision_max_wh_ratio is not None and width > height * args.vision_max_wh_ratio:
        width = height * args.vision_max_wh_ratio
    if 0 < width * height < args.vision_min_pixels:
        ratio = (args.vision_min_pixels / (width * height)) ** 0.5
        width = int(width * ratio)
        height = int(height * ratio)
    best_width = math.ceil(width / p) * p
    best_height = math.ceil(height / p) * p
    n_llm_h, n_llm_w, best_height, best_width = safe_resize(
        height,
        width,
        best_height,
        best_width,
        p,
        args.vision_downsample_ratio,
        args.vision_max_n_token,
    )
    n_vit_h, n_vit_w = best_height // p, best_width // p
    src_w, src_h = image.size
    if args.vision_max_wh_ratio is not None and src_w >= args.vision_max_wh_ratio * src_h:
        image = image.resize((best_width, best_height))
    else:
        image = ImageOps.pad(image, (best_width, best_height), color=(127, 127, 127))
    try:
        import numpy as np

        x = torch.from_numpy(np.asarray(image, dtype=np.float32)).permute(2, 0, 1) / 255
    except ImportError:
        pixels = (
            image.get_flattened_data()
            if hasattr(image, "get_flattened_data")
            else image.getdata()
        )
        x = torch.tensor(pixels, dtype=torch.float32)
        x = x.view(image.height, image.width, 3).permute(2, 0, 1) / 255
    x = ((x - 0.5) / 0.5).to(torch.bfloat16)
    patches = (
        x.reshape(3, n_vit_h, p, n_vit_w, p)
        .permute(1, 3, 0, 2, 4)
        .reshape(n_vit_h * n_vit_w, 3, p, p)
    )
    return patches, n_vit_h, n_vit_w, n_llm_h, n_llm_w


def load_image(record, args):
    """Load and transform one image record into ViT patches."""
    Image, _ImageOps = _pil()
    with Image.open(io.BytesIO(load_image_bytes(record))) as source:
        image = source.convert("RGB")
        return pil_to_patches(image, args)


def build_image_block(n_llm_h: int, n_llm_w: int, start_pos: int):
    """Builds the N-layout token types (final order) and the aligner-row order for IMAGE slots."""
    import torch

    compress_pad = COMPRESS_PAD_TO - 1 - start_pos % COMPRESS_PAD_TO
    pad_h = n_llm_h % 2
    rows = n_llm_h + pad_h
    row_len = n_llm_w + 1
    pad_last = rows // 2 * row_len % 2 * 2
    types = torch.tensor(
        ([IMAGE] * n_llm_w + [IMAGE_NEW_LINE]) * n_llm_h + [IMAGE_PAD] * (row_len * pad_h),
        dtype=torch.int64,
    )
    order = torch.arange(rows * row_len).view(rows // 2, 2, row_len).transpose(1, 2).reshape(-1)
    image_idx = torch.full((rows * row_len,), -1, dtype=torch.int64)
    image_idx.view(rows, row_len)[:n_llm_h, :n_llm_w] = torch.arange(n_llm_h * n_llm_w).view(
        n_llm_h, n_llm_w
    )
    perm = image_idx[order]
    perm = perm[perm >= 0]
    types = torch.cat(
        [
            torch.full((compress_pad,), IMAGE_PAD, dtype=torch.int64),
            torch.tensor([IMAGE_START]),
            types[order],
            torch.full((pad_last,), IMAGE_PAD, dtype=torch.int64),
            torch.tensor([IMAGE_END]),
        ]
    )
    return types, perm