"""Pure-CPU unit tests for _chunk_mm_metadata: lock down the boundary semantics of bisecting image-range intersections, byte-range slicing, and patch-dimension slicing.

No real ckpt / GPU compute is needed: only CPU-tensor metadata dicts are built and passed to pure functions; importing deepinfer still needs _C (a GPU machine).
"""

from typing import Any

import pytest
import torch

from deepinfer.ve.prepare_tensor_input_vl import _chunk_mm_metadata, _prepare_raw_image_transform_inputs

# The raw byte count / patch count per image is fixed, which makes the rebased byte and patch ranges easy to assert.
RAW_BYTES_PER_IMG = 12
PATCHES_PER_IMG = 3


def _md(spans: list[tuple[int, int]], patches_per_img: int = PATCHES_PER_IMG) -> dict[str, Any]:
    """Build a metadata dict in the mapper's output shape from per-image (start, count).

    spans: [(start, count)] is where image_starts/image_lens come from; each image has RAW_BYTES_PER_IMG raw RGB bytes and
    patches_per_img patches. The contents of images_raw are global byte indices (so the sliced ranges can be asserted directly).
    """
    starts = torch.tensor([s for s, _ in spans], dtype=torch.int32)
    lens = torch.tensor([c for _, c in spans], dtype=torch.int32)
    n = len(spans)
    cu_patches = torch.zeros(n + 1, dtype=torch.int64)
    cu_raw = torch.zeros(n + 1, dtype=torch.int64)
    for i in range(n):
        cu_patches[i + 1] = cu_patches[i] + patches_per_img
        cu_raw[i + 1] = cu_raw[i] + RAW_BYTES_PER_IMG
    total_patch = int(cu_patches[-1])
    total_raw = int(cu_raw[-1])
    return {
        "images_raw": torch.arange(total_raw, dtype=torch.uint8) if total_raw else torch.zeros(0, dtype=torch.uint8),
        "images_input_size": torch.full((n, 2), 4, dtype=torch.int64) if n else torch.zeros(0, 2, dtype=torch.int64),
        "cu_raw_bytes": cu_raw,
        "image_starts": starts,
        "image_lens": lens,
        "cu_image_patch_lens": cu_patches,
        "images_spatial_crop": torch.arange(n * 2).reshape(n, 2) if n else torch.zeros(0, 2, dtype=torch.int64),
        "images_global_attn_pos_id": torch.arange(total_patch * 2).reshape(total_patch, 2),
        "images_seq_mask": None,
        "mm_token_type": None,
    }


def _hit_range(spans: list[tuple[int, int]], start_idx: int, end_idx: int) -> tuple[int, int] | None:
    """Reference for the old min/max semantics: images intersecting the chunk are [first, last+1); None when there are none."""
    hit = [i for i, (s, c) in enumerate(spans) if s < end_idx and s + c > start_idx]
    return (hit[0], hit[-1] + 1) if hit else None


def _assert_chunk(spans: list[tuple[int, int]], start_idx: int, end_idx: int, patches_per_img: int = PATCHES_PER_IMG) -> None:
    md = _chunk_mm_metadata(_md(spans, patches_per_img), start_idx, end_idx)
    rng = _hit_range(spans, start_idx, end_idx)
    if rng is None:
        # image-free chunk: all image-related keys are set to None (downstream skips image encoding as "no images")
        for k in (
            "images_raw",
            "images_input_size",
            "cu_raw_bytes",
            "images_spatial_crop",
            "images_global_attn_pos_id",
            "cu_image_patch_lens",
        ):
            assert md[k] is None, k
        return
    si, ei = rng
    # images_raw is sliced by byte range: this chunk's raw bytes = [si*RAW, ei*RAW)
    assert md["images_raw"].tolist() == list(range(si * RAW_BYTES_PER_IMG, ei * RAW_BYTES_PER_IMG))
    assert md["images_input_size"].shape[0] == ei - si
    # cu_raw_bytes is rebased to this chunk's start and accumulated per image
    assert md["cu_raw_bytes"].tolist() == [i * RAW_BYTES_PER_IMG for i in range(ei - si + 1)]
    # pos_id is sliced by patch range; cu_image_patch_lens is rebased to start at 0
    assert md["images_global_attn_pos_id"].shape[0] == (ei - si) * patches_per_img
    assert md["cu_image_patch_lens"][0].item() == 0
    assert md["cu_image_patch_lens"][-1].item() == (ei - si) * patches_per_img
    assert md["images_spatial_crop"].shape[0] == ei - si


def test_chunk_at_image_end_excludes_image() -> None:
    # image0=[0,3), image1=[3,6); the chunk ends exactly at image0's end, end_idx==3 => only image0 is included
    spans = [(0, 3), (3, 3)]
    _assert_chunk(spans, 0, 3)
    md = _chunk_mm_metadata(_md(spans), 0, 3)
    assert md["images_raw"].shape[0] == RAW_BYTES_PER_IMG  # only image0's bytes
    assert md["cu_image_patch_lens"][-1].item() == PATCHES_PER_IMG  # only image0's patches


def test_chunk_at_image_start_includes_image() -> None:
    # the chunk starts at image1's start, start_idx==3 => image1 is included (image0 is not)
    spans = [(0, 3), (3, 3)]
    md = _chunk_mm_metadata(_md(spans), 3, 6)
    assert md["images_raw"].tolist() == list(range(RAW_BYTES_PER_IMG, 2 * RAW_BYTES_PER_IMG))
    assert md["images_spatial_crop"].shape[0] == 1


def test_empty_chunk_has_no_image() -> None:
    spans = [(0, 3), (3, 3)]
    md = _chunk_mm_metadata(_md(spans), 3, 3)  # start==end
    assert md["images_raw"] is None
    assert md["cu_image_patch_lens"] is None


def test_multi_image_partial_overlap() -> None:
    # images [2,4), [5,7), [8,11); chunk [3,9) intersects all three
    spans = [(2, 2), (5, 2), (8, 3)]
    for s in range(12):
        for e in range(s + 1, 13):
            _assert_chunk(spans, s, e, patches_per_img=4)


def test_no_image_returns_early() -> None:
    md = _chunk_mm_metadata(_md([]), 0, 4)
    # no images: image_starts is an empty tensor and intersect is empty -> all image-related keys are set to None
    assert md["images_raw"] is None
    assert md["images_spatial_crop"] is None


def test_no_raw_bytes_drops_image_metadata() -> None:
    """When this seq has no raw bytes (every image was cache-hit and pre-split), intersecting a zero-placeholder image must not crash:
    images_input_size/cu_raw_bytes are None -> the whole image key group is set to None and downstream treats it as image-free."""
    metadata = _md([(0, 3)])
    metadata["images_input_size"] = None
    metadata["images_raw"] = None
    metadata["cu_raw_bytes"] = None
    md = _chunk_mm_metadata(metadata, 0, 3)
    for k in ("images_raw", "images_input_size", "cu_raw_bytes", "images_spatial_crop", "cu_image_patch_lens"):
        assert md[k] is None, k


def test_does_not_mutate_input() -> None:
    """Shallow-copy semantics: tensors produced by slicing must not write back into the original metadata."""
    metadata = _md([(0, 3), (3, 3)])
    orig_raw = metadata["images_raw"]
    orig_cu_raw = metadata["cu_raw_bytes"]
    orig_cu_patch = metadata["cu_image_patch_lens"]
    _chunk_mm_metadata(metadata, 3, 6)
    assert metadata["images_raw"] is orig_raw
    assert metadata["cu_raw_bytes"] is orig_cu_raw
    assert metadata["cu_image_patch_lens"] is orig_cu_patch
    assert metadata["cu_raw_bytes"].tolist() == [0, RAW_BYTES_PER_IMG, 2 * RAW_BYTES_PER_IMG]  # unaffected by rebase


@pytest.mark.parametrize("sizes", [[(1302, 1302)] * 424, [(65536, 32768), (14, 14)]])
@pytest.mark.parametrize("dtype", [torch.int32, torch.int64])
def test_prepare_raw_image_inputs_int64(sizes: list[tuple[int, int]], dtype: torch.dtype) -> None:
    """Realistic multi-image overflow plus a single-image multiply-before-cast regression.

    Only the raw buffer is meta: preparation never reads pixels, so this exercises
    the real CPU arithmetic without allocating several GiB of image data.
    """
    byte_lens = [w * h * 3 for w, h in sizes]
    raw = torch.empty(sum(byte_lens), dtype=torch.uint8, device="meta")
    metadata = {"images_raw": raw, "images_input_size": torch.tensor(sizes, dtype=dtype)}
    _prepare_raw_image_transform_inputs(metadata, total_image_patches=123)
    expected = [0]
    for n in byte_lens:
        expected.append(expected[-1] + n)
    assert expected[-1] > 2**31
    assert metadata["cu_raw_bytes"].dtype == torch.int64
    assert metadata["cu_raw_bytes"].tolist() == expected
    assert metadata["image_raw_sizes"].dtype == torch.int32
    assert metadata["image_raw_sizes"].tolist() == [list(s) for s in sizes]
    assert metadata["images_raw"] is raw
    assert metadata["total_image_patches"] == 123
    assert "images_input_size" not in metadata


def test_chunk_raw_bytes_above_int32() -> None:
    metadata = _md([(0, 3), (3, 6)])
    raw_start = 2**31 + 12
    metadata["images_raw"] = torch.empty(raw_start + 12, dtype=torch.uint8, device="meta")
    metadata["cu_raw_bytes"] = torch.tensor([0, raw_start, raw_start + 12], dtype=torch.int64)
    chunk = _chunk_mm_metadata(metadata, 3, 9)
    assert chunk["images_raw"].numel() == 12
    assert chunk["images_raw"].storage_offset() == raw_start
    assert chunk["cu_raw_bytes"].dtype == torch.int64
    assert chunk["cu_raw_bytes"].tolist() == [0, 12]


def test_prepare_empty_raw_image_inputs() -> None:
    metadata = {"images_raw": None, "images_input_size": None}
    _prepare_raw_image_transform_inputs(metadata, total_image_patches=0)
    assert metadata["images_raw"] is None
    assert metadata["image_raw_sizes"] is None
    assert metadata["cu_raw_bytes"] is None
    assert metadata["total_image_patches"] == -1