# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import random
from typing import Any, NamedTuple, Optional, cast

import numpy as np
import pytest
from transformers import AutoTokenizer, PreTrainedTokenizerBase

from vllm.benchmarks.datasets import (RandomDataset, RandomMultiModalDataset,
                                      SampleRequest)


@pytest.fixture(scope="session")
def hf_tokenizer() -> PreTrainedTokenizerBase:
    # Use a small, commonly available tokenizer
    return AutoTokenizer.from_pretrained("gpt2")


class Params(NamedTuple):
    num_requests: int
    prefix_len: int
    range_ratio: float
    input_len: int
    output_len: int


@pytest.fixture(scope="session")
def random_dataset_params() -> Params:
    return Params(num_requests=16,
                  prefix_len=7,
                  range_ratio=0.3,
                  input_len=50,
                  output_len=20)


def _fingerprint_sample(req: SampleRequest) -> tuple[str, int, int]:
    """Project a SampleRequest into a comparable tuple."""
    return (req.prompt, req.prompt_len, req.expected_output_len)


def _collect_samples(dataset: RandomDataset,
                    tokenizer: PreTrainedTokenizerBase,
                     num_requests: int = 16,
                     prefix_len: int = 7,
                     range_ratio: float = 0.3,
                     input_len: int = 50,
                     output_len: int = 20) -> list[tuple[str, int, int]]:
    samples = dataset.sample(
        tokenizer=tokenizer,
        num_requests=num_requests,
        prefix_len=prefix_len,
        range_ratio=range_ratio,
        input_len=input_len,
        output_len=output_len,
    )
    return [_fingerprint_sample(s) for s in samples]


@pytest.mark.benchmark
def test_random_dataset_same_seed(
        hf_tokenizer: PreTrainedTokenizerBase,
        random_dataset_params: Params) -> None:
    """Same seed should yield identical outputs, even if global RNGs change.

    This guards against accidental reliance on Python's random or np.random
    in RandomDataset after moving to numpy.default_rng.
    """
    p = random_dataset_params
    common_seed = 123
    dataset_a = RandomDataset(random_seed=common_seed)
    dataset_b = RandomDataset(random_seed=common_seed)
    a = _collect_samples(dataset_a,
                         hf_tokenizer,
                         num_requests=p.num_requests,
                         prefix_len=p.prefix_len,
                         range_ratio=p.range_ratio,
                         input_len=p.input_len,
                         output_len=p.output_len)

    # Perturb global RNG state to ensure isolation
    random.seed(999)
    _ = [random.random() for _ in range(100)]
    np.random.seed(888)
    _ = [np.random.random() for _ in range(100)]

    b = _collect_samples(dataset_b,
                         hf_tokenizer,
                         num_requests=p.num_requests,
                         prefix_len=p.prefix_len,
                         range_ratio=p.range_ratio,
                         input_len=p.input_len,
                         output_len=p.output_len)
    assert a == b

@pytest.mark.benchmark
def test_random_dataset_different_seeds(
        hf_tokenizer: PreTrainedTokenizerBase,
        random_dataset_params: Params) -> None:
    """Different seeds should change outputs with overwhelming likelihood."""
    p = random_dataset_params
    seed_a = 0
    dataset_a = RandomDataset(random_seed=seed_a)
    a = _collect_samples(dataset_a,
                         hf_tokenizer,
                         num_requests=p.num_requests,
                         prefix_len=p.prefix_len,
                         range_ratio=p.range_ratio,
                         input_len=p.input_len,
                         output_len=p.output_len)

    seed_b = 999
    dataset_b = RandomDataset(random_seed=seed_b)
    # Perturb global RNG with same seed as dataset_a to ensure isolation
    random.seed(seed_a)
    np.random.seed(seed_a)
    b = _collect_samples(dataset_b,
                         hf_tokenizer,
                         num_requests=p.num_requests,
                         prefix_len=p.prefix_len,
                         range_ratio=p.range_ratio,
                         input_len=p.input_len,
                         output_len=p.output_len)
    assert a != b


# -----------------------------
# RandomMultiModalDataset tests
# -----------------------------

def _mm_fingerprint_sample(
    req: SampleRequest,
) -> tuple[str, int, int, int, list[str]]:
    """Create a compact fingerprint for multimodal samples.

    Includes:
    - prompt string
    - prompt_len
    - expected_output_len
    - count of multimodal items
    - per-item type and URL prefix (e.g., 'data:image/jpeg;base64,')
    """
    items = req.multi_modal_data or []
    item_prefixes: list[str] = []
    for it in items:
        if isinstance(it, dict) and it.get("type") == "image_url":
            url = it.get("image_url", {}).get("url", "")
            # Only keep a short identifying prefix to avoid huge strings
            item_prefixes.append(f"image:{url[:22]}")
        elif isinstance(it, dict) and it.get("type") == "video_url":
            url = it.get("video_url", {}).get("url", "")
            item_prefixes.append(f"video:{url[:22]}")
        else:
            item_prefixes.append("unknown:")
    return (req.prompt, req.prompt_len, req.expected_output_len, len(items),
            item_prefixes)


def _collect_mm_samples(
    dataset: RandomMultiModalDataset,
    tokenizer: PreTrainedTokenizerBase,
    *,
    num_requests: int = 8,
    prefix_len: int = 3,
    range_ratio: float = 0.0,
    input_len: int = 20,
    output_len: int = 5,
    base_items_per_request: int = 2,
    num_mm_items_range_ratio: float = 0.0,
    limit_mm_per_prompt: Optional[dict[str, int]] = None,
    bucket_config: Optional[dict[tuple[int, int, int], float]] = None,
    enable_multimodal_chat: bool = False,
) -> list[SampleRequest]:
    if limit_mm_per_prompt is None:
        limit_mm_per_prompt = {"image": 5, "video": 0}
    if bucket_config is None:
        bucket_config = {(32, 32, 1): 0.5, (52, 64, 1): 0.5}
    return dataset.sample(
        tokenizer=tokenizer,
        num_requests=num_requests,
        prefix_len=prefix_len,
        range_ratio=range_ratio,
        input_len=input_len,
        output_len=output_len,
        base_items_per_request=base_items_per_request,
        num_mm_items_range_ratio=num_mm_items_range_ratio,
        limit_mm_per_prompt=limit_mm_per_prompt,
        bucket_config=bucket_config,
        enable_multimodal_chat=enable_multimodal_chat,
    )


@pytest.mark.benchmark
def test_random_mm_same_seed(hf_tokenizer: PreTrainedTokenizerBase) -> None:
    seed = 42
    ds_a = RandomMultiModalDataset(random_seed=seed)
    ds_b = RandomMultiModalDataset(random_seed=seed)
    a = _collect_mm_samples(ds_a, hf_tokenizer)
    b = _collect_mm_samples(ds_b, hf_tokenizer)
    fa = [_mm_fingerprint_sample(s) for s in a]
    fb = [_mm_fingerprint_sample(s) for s in b]
    assert fa == fb


@pytest.mark.benchmark
def test_random_mm_different_seeds(
    hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
    ds_a = RandomMultiModalDataset(random_seed=0)
    ds_b = RandomMultiModalDataset(random_seed=999)
    a = _collect_mm_samples(ds_a, hf_tokenizer)
    b = _collect_mm_samples(ds_b, hf_tokenizer)
    fa = [_mm_fingerprint_sample(s) for s in a]
    fb = [_mm_fingerprint_sample(s) for s in b]
    assert fa != fb

@pytest.mark.benchmark
def test_random_mm_respects_limits(
    hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
    ds = RandomMultiModalDataset(random_seed=0)
    # Requesting 3 items with a per-prompt limit of 1 should error per current
    # design (dataset refuses to silently clamp below the requested baseline).
    with pytest.raises(ValueError):
        _collect_mm_samples(
            ds,
            hf_tokenizer,
            num_requests=12,
            base_items_per_request=3,
            num_mm_items_range_ratio=0.0,
            limit_mm_per_prompt={"image": 1, "video": 0},
            bucket_config={(32, 32, 1): 1.0},
        )


@pytest.mark.benchmark
def test_random_mm_zero_prob_entries_are_removed(
    hf_tokenizer: PreTrainedTokenizerBase,
) -> None:
    ds = RandomMultiModalDataset(random_seed=0)
    # Second bucket has zero probability and should be ignored after
    # normalization
    samples = _collect_mm_samples(
        ds,
        hf_tokenizer,
        num_requests=6,
        base_items_per_request=2,
        num_mm_items_range_ratio=0.0,
        limit_mm_per_prompt={"image": 10, "video": 0},
        bucket_config={(32, 32, 1): 1.0, (52, 64, 1): 0.0},
    )
    for s in samples:
        assert isinstance(s.multi_modal_data, list)
        typed_mm = cast(list[dict[str, Any]], s.multi_modal_data)
        for it in typed_mm:
            assert it.get("type") == "image_url"


@pytest.mark.benchmark
def test_random_mm_zero_items(hf_tokenizer: PreTrainedTokenizerBase) -> None:
    ds = RandomMultiModalDataset(random_seed=0)
    samples = _collect_mm_samples(
        ds,
        hf_tokenizer,
        num_requests=5,
        base_items_per_request=0,
        num_mm_items_range_ratio=0.0,
        limit_mm_per_prompt={"image": 5, "video": 0},
        bucket_config={(32, 32, 1): 1.0},
    )
    for s in samples:
        assert s.multi_modal_data == []

@pytest.mark.benchmark
def test_random_mm_num_items_per_prompt(
    hf_tokenizer: PreTrainedTokenizerBase) -> None:
    ds = RandomMultiModalDataset(random_seed=0)
    # Fixed number of images per prompt
    # set num_mm_items_range_ratio to 0.0
    # TODO: modify video values when video sampling is implemented
    samples_fixed_items = _collect_mm_samples(
        ds,
        hf_tokenizer,
        num_requests=5,
        base_items_per_request=3,
        num_mm_items_range_ratio=0.0,
        limit_mm_per_prompt={"image": 3, "video": 0},
        bucket_config={(32, 32, 1): 1.0},
    )
    # Must have 5 requests each with 3 mm items per prompt
    assert len(samples_fixed_items) == 5
    for s in samples_fixed_items:
        mm_data = cast(list[dict[str, Any]], s.multi_modal_data)
        assert len(mm_data) == 3
        for it in mm_data:
            assert it.get("type") == "image_url"


@pytest.mark.benchmark
def test_random_mm_bucket_config_not_mutated(
    hf_tokenizer: PreTrainedTokenizerBase,
) -> None:

    ds = RandomMultiModalDataset(random_seed=0)
    # This bucket config is not normalized to sum to 1
    # and has more buckets than requested images
    original = {(32, 32, 1): 0.2, (52, 64, 1): 6, (25, 64, 1): 3}
    # Keep a snapshot to compare after sampling
    snapshot = dict(original)

    _ = _collect_mm_samples(
        ds,
        hf_tokenizer,
        num_requests=4,
        base_items_per_request=1,
        num_mm_items_range_ratio=0.0,
        limit_mm_per_prompt={"image": 1, "video": 0},
        bucket_config=original,
    )

    # Ensure the original dict content is unchanged
    assert original == snapshot


    # Vary number of mm items per prompt
    # set num_mm_items_range_ratio to 0.5
    samples_varying_items = _collect_mm_samples(
        ds,
        hf_tokenizer,
        num_requests=5,
        base_items_per_request=2,
        num_mm_items_range_ratio=0.5,
        limit_mm_per_prompt={"image": 4, "video": 0},
        bucket_config={(32, 32, 1): 1.0},
    )
    # Must have 5 requests each with less than 4 mm items per prompt
    # but at least 1 mm item per prompt
    assert len(samples_varying_items) == 5
    for s in samples_varying_items:
        mm_data = cast(list[dict[str, Any]], s.multi_modal_data)
        assert len(mm_data) <= 4
        assert len(mm_data) >= 1
        for it in mm_data:
            assert it.get("type") == "image_url"
