# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

import pytest

from vllm_omni.diffusion.diffusion_kv.metadata import (
    DiffusionKVContextMetadata,
    DiffusionKVMetadata,
    DiffusionKVSequenceMetadata,
)
from vllm_omni.diffusion.diffusion_kv.request import DiffusionKVRequest
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched.base_scheduler import BaseScheduler
from vllm_omni.diffusion.sched.interface import CachedRequestData, DiffusionSchedulerOutput
from vllm_omni.inputs.data import OmniDiffusionSamplingParams

pytestmark = [pytest.mark.diffusion, pytest.mark.core_model, pytest.mark.cpu]


def test_native_request_metadata_round_trips_and_stays_optional() -> None:
    params = {
        "transfer_id": "xfer-req-0",
        "do_remote_decode": True,
        "do_remote_prefill": False,
        "remote_engine_id": "dit-engine-0",
    }
    kv_request = DiffusionKVRequest(
        "req-0/diffusion-kv/0",
        sequence_id=0,
        prefix_len=4,
        target_len=2,
        seq_len=6,
        kv_transfer_params=params,
        prompt_token_ids=[0] * 6,
    )
    public_request = OmniDiffusionRequest(
        prompt="prompt",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="req-0",
        kv_transfer_params=params,
    )
    scheduler_output = DiffusionSchedulerOutput(
        step_id=0,
        scheduled_new_reqs=[],
        scheduled_cached_reqs=CachedRequestData.make_empty(),
        finished_req_ids=set(),
        num_running_reqs=0,
        num_waiting_reqs=0,
    )

    assert kv_request.kv_transfer_params is params
    assert kv_request.prompt_token_ids == [0] * 6
    assert public_request.kv_transfer_params is params
    assert scheduler_output.kv_connector_metadata is None


class _MetadataScheduler(BaseScheduler):
    def update_from_output(self, sched_output, output) -> set[str]:
        del sched_output, output
        return set()


def test_scheduler_copies_opaque_params_to_sequence_requests() -> None:
    params = {"transfer_id": "xfer-req-0", "do_remote_prefill": True, "num_transfer_tokens": 3}
    sequence = DiffusionKVRequest(
        "req-0/diffusion-kv/0",
        sequence_id=0,
        prefix_len=4,
        target_len=2,
        seq_len=6,
        prompt_token_ids=[10, 11, 12, 13],
    )
    request = OmniDiffusionRequest(
        prompt="prompt",
        sampling_params=OmniDiffusionSamplingParams(num_inference_steps=1),
        request_id="req-0",
        diffusion_kv_requests=(sequence,),
        kv_transfer_params=params,
    )
    scheduler = _MetadataScheduler()
    scheduler._diffusion_kv_manager = object()

    state = scheduler._make_request_state("req-0", request)

    assert state.diffusion_kv_requests[0].kv_transfer_params == params
    assert state.diffusion_kv_requests[0].kv_transfer_params is not params
    assert sequence.prompt_token_ids == [10, 11, 12]
    assert sequence.num_prompt_tokens == 3
    assert sequence.seq_len == 6


def test_diffusion_kv_metadata_uses_native_cache_group_block_ids() -> None:
    context = DiffusionKVContextMetadata(
        context_id="text",
        cache_role="cross_attention",
        num_tokens=3,
        block_ids=([7], [11]),
    )
    sequence = DiffusionKVSequenceMetadata(
        sequence_id=1,
        prefix_len=4,
        target_len=2,
        seq_len=8,
        block_ids=([1, 2], [5, 6]),
        context_ids=(context.context_id,),
    )
    metadata = DiffusionKVMetadata(
        request_id="req-0",
        allocation_generation=3,
        sequences=(sequence,),
        contexts=(context,),
    )

    assert metadata.sequences[0].block_ids == ([1, 2], [5, 6])
    assert metadata.contexts[0].block_ids == ([7], [11])


def test_cfg_sequences_can_share_one_request_context() -> None:
    shared_context = DiffusionKVContextMetadata(
        context_id="shared-text",
        cache_role="cross_attention",
        num_tokens=3,
        block_ids=([7], [11]),
    )
    sequences = tuple(
        DiffusionKVSequenceMetadata(
            sequence_id=sequence_id,
            prefix_len=4,
            target_len=2,
            seq_len=8,
            block_ids=([sequence_id + 1],),
            context_ids=(shared_context.context_id,),
        )
        for sequence_id in range(2)
    )

    metadata = DiffusionKVMetadata(
        request_id="req-shared-context",
        allocation_generation=1,
        sequences=sequences,
        contexts=(shared_context,),
    )

    assert metadata.contexts == (shared_context,)
    assert metadata.sequences[0].context_ids == (shared_context.context_id,)
    assert metadata.sequences[1].context_ids == (shared_context.context_id,)


def test_cfg_sequences_can_reference_branch_specific_contexts() -> None:
    conditional_context = DiffusionKVContextMetadata(
        context_id="conditional-text",
        cache_role="cross_attention",
        num_tokens=3,
        block_ids=([7], [11]),
    )
    unconditional_context = DiffusionKVContextMetadata(
        context_id="unconditional-text",
        cache_role="cross_attention",
        num_tokens=3,
        block_ids=([8], [12]),
    )
    metadata = DiffusionKVMetadata(
        request_id="req-branch-contexts",
        allocation_generation=1,
        sequences=(
            DiffusionKVSequenceMetadata(
                sequence_id=0,
                prefix_len=4,
                target_len=2,
                seq_len=8,
                block_ids=([1],),
                context_ids=(conditional_context.context_id,),
            ),
            DiffusionKVSequenceMetadata(
                sequence_id=1,
                prefix_len=4,
                target_len=2,
                seq_len=8,
                block_ids=([2],),
                context_ids=(unconditional_context.context_id,),
            ),
        ),
        contexts=(conditional_context, unconditional_context),
    )

    assert metadata.sequences[0].context_ids == (conditional_context.context_id,)
    assert metadata.sequences[1].context_ids == (unconditional_context.context_id,)