import pytest
import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import SamplingParams
from vllm.v1.engine import EngineCoreRequest

from vllm_omni.distributed.omni_coordinator import ReplicaInfo, ReplicaStatus
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine, StageRuntimeInfo
from vllm_omni.engine.serialization import deserialize_additional_information
from vllm_omni.engine.stage_pool import StagePool
from vllm_omni.model_executor.stage_input_processors.bagel import ExpandedPrompt

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


def _make_engine_core_request(request_id: str = "req-1") -> EngineCoreRequest:
    return EngineCoreRequest(
        request_id=request_id,
        prompt_token_ids=[1, 1, 1],
        mm_features=None,
        sampling_params=SamplingParams(max_tokens=8),
        pooling_params=None,
        arrival_time=0.0,
        lora_request=None,
        cache_salt=None,
        data_parallel_rank=None,
    )


def test_build_add_request_message_preserves_additional_information(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("speech",)

    input_processor = mocker.Mock()
    input_processor.process_inputs.return_value = _make_engine_core_request()
    engine.input_processor = input_processor

    output_processor = mocker.Mock()
    engine.output_processors = [output_processor]

    prompt = {
        "prompt_token_ids": [1, 1, 1],
        "additional_information": {
            "text": ["hello world"],
            "speaker": ["vivian"],
        },
    }

    msg = engine._build_add_request_message(
        request_id="req-1",
        prompt=prompt,
        sampling_params_list=[params],
        final_stage_id=0,
        arrival_time=0.0,
    )

    request = msg.prompt
    assert isinstance(request, OmniEngineCoreRequest)
    assert request.external_req_id == "req-1"
    assert request.additional_information is not None
    assert request.additional_information.entries["text"].list_data == ["hello world"]
    assert request.additional_information.entries["speaker"].list_data == ["vivian"]
    output_processor.add_request.assert_not_called()


def test_build_add_request_message_preserves_model_intermediate_buffer(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("speech",)

    input_processor = mocker.Mock()
    input_processor.process_inputs.return_value = _make_engine_core_request()
    engine.input_processor = input_processor

    output_processor = mocker.Mock()
    engine.output_processors = [output_processor]

    hidden = torch.arange(6, dtype=torch.float32).reshape(2, 3)
    prompt = {
        "prompt_token_ids": [1, 1, 1],
        "model_intermediate_buffer": {
            "ids": {"tts": [11, 12]},
            "hidden_states": {"tts": hidden},
        },
    }

    msg = engine._build_add_request_message(
        request_id="req-1",
        prompt=prompt,
        sampling_params_list=[params],
        final_stage_id=0,
        arrival_time=0.0,
    )

    request = msg.prompt
    assert isinstance(request, OmniEngineCoreRequest)
    assert request.additional_information is not None
    assert request.additional_information.entries["global_request_id"].list_data == ["req-1"]
    assert request.additional_information.entries["omni_final_stage_id"].scalar_data == 0
    assert isinstance(request.model_intermediate_buffer, dict)
    info = request.model_intermediate_buffer
    assert info["ids"]["tts"] == [11, 12]
    assert torch.equal(info["hidden_states"]["tts"], hidden)


def test_cfg_companion_suppresses_payload_but_forces_kv_transfer(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.prompt_expand_func = lambda *_args: [
        ExpandedPrompt(
            prompt={"prompt": "negative"},
            role="cfg_text",
            request_id_suffix="__cfg_text",
        )
    ]
    engine.supported_tasks = ("generate",)
    engine.input_processor = mocker.Mock()
    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
    engine.request_queue = mocker.Mock()

    engine._enqueue_cfg_companions(
        parent_id="req",
        original_prompt={"prompt": "positive"},
        stage0_params=params,
        sampling_params_list=[params],
    )

    message = engine.request_queue.sync_q.put.call_args.args[0]
    metadata = deserialize_additional_information(message.prompt.additional_information)
    assert metadata["omni_final_stage_id"] == 0
    assert metadata["omni_force_kv_transfer"] is True


def test_build_add_request_message_with_resumable_streaming(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("generate",)

    input_processor = mocker.Mock()
    input_processor.process_inputs.return_value = _make_engine_core_request()
    engine.input_processor = input_processor

    output_processor = mocker.Mock()
    engine.output_processors = [output_processor]

    msg = engine._build_add_request_message(
        request_id="req-stream",
        prompt={"prompt_token_ids": [1, 2, 3]},
        sampling_params_list=[params],
        final_stage_id=0,
        resumable=True,
        message_type="streaming_update",
    )

    assert msg.type == "streaming_update"
    input_processor.process_inputs.assert_called_once()
    assert input_processor.process_inputs.call_args.kwargs["resumable"] is True


class _FakeStageClient:
    stage_type = "llm"
    final_output = False

    def __init__(self, input_address: str | None = None):
        if input_address is not None:
            self.client_addresses = {"input_address": input_address}


class _FakeHub:
    def __init__(self, replicas: list[ReplicaInfo]):
        self._replicas = replicas

    def get_replicas_for_stage(self, stage_id: int):
        return type(
            "ReplicaList",
            (),
            {"replicas": [rep for rep in self._replicas if rep.stage_id == stage_id]},
        )()


class _RoundRobinLB:
    def __init__(self):
        self._next = 0

    def select(self, task, replicas):  # noqa: ARG002
        idx = self._next % len(replicas)
        self._next += 1
        return idx


def _replica(input_addr: str) -> ReplicaInfo:
    return ReplicaInfo(
        input_addr=input_addr,
        output_addr=input_addr.replace("input", "output"),
        stage_id=0,
        status=ReplicaStatus.UP,
        queue_length=0,
        last_heartbeat=0.0,
        registered_at=0.0,
    )


def test_build_add_request_message_scopes_mm_uuids_to_selected_stage0_replica(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.model = "test-model"
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("generate",)
    engine.stage_pools = [StagePool(0, [_FakeStageClient(), _FakeStageClient()])]

    seen_uuids: list[str] = []

    def process_inputs(**kwargs):
        prompt = kwargs["prompt"]
        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
        return _make_engine_core_request(kwargs["request_id"])

    input_processor = mocker.Mock()
    input_processor.process_inputs.side_effect = process_inputs
    engine.input_processor = input_processor

    for request_id in ("req-1", "req-2"):
        engine._build_add_request_message(
            request_id=request_id,
            prompt={
                "prompt": "describe",
                "multi_modal_data": {"image": "same-image"},
            },
            sampling_params_list=[params],
            final_stage_id=0,
        )

    assert seen_uuids[0].startswith("stage0:rep0:")
    assert seen_uuids[1].startswith("stage0:rep1:")
    assert seen_uuids[0].removeprefix("stage0:rep0:") == seen_uuids[1].removeprefix("stage0:rep1:")


@pytest.mark.asyncio
async def test_build_add_request_message_scopes_mm_uuids_to_distributed_stage0_replica(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.model = "test-model"
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("generate",)

    addr0 = "tcp://host-a:1000/input"
    addr1 = "tcp://host-b:1000/input"
    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
    stage_pool.attach_hub(_FakeHub([_replica(addr0), _replica(addr1)]))
    stage_pool.attach_load_balancer(_RoundRobinLB())
    engine.stage_pools = [stage_pool]

    seen_uuids: list[str] = []

    def process_inputs(**kwargs):
        prompt = kwargs["prompt"]
        seen_uuids.append(prompt["multi_modal_uuids"]["image"][0])
        return _make_engine_core_request(kwargs["request_id"])

    input_processor = mocker.Mock()
    input_processor.process_inputs.side_effect = process_inputs
    engine.input_processor = input_processor

    for request_id in ("req-1", "req-2"):
        engine._build_add_request_message(
            request_id=request_id,
            prompt={
                "prompt": "describe",
                "multi_modal_data": {"image": "same-image"},
            },
            sampling_params_list=[params],
            final_stage_id=0,
        )

    assert seen_uuids[0].startswith("stage0:rep0:")
    assert seen_uuids[1].startswith("stage0:rep1:")
    assert stage_pool.get_bound_replica_id("req-1") == 0
    assert stage_pool.get_bound_replica_id("req-2") == 1
    assert await stage_pool.pick("req-1") == 0
    assert await stage_pool.pick("req-2") == 1


def test_build_add_request_message_skips_distributed_mm_scope_when_no_replica(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.model = "test-model"
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("generate",)

    addr0 = "tcp://host-a:1000/input"
    addr1 = "tcp://host-b:1000/input"
    stage_pool = StagePool(0, [_FakeStageClient(addr0), _FakeStageClient(addr1)])
    stage_pool.attach_hub(_FakeHub([]))
    stage_pool.attach_load_balancer(_RoundRobinLB())
    engine.stage_pools = [stage_pool]

    seen_prompt: dict | None = None

    def process_inputs(**kwargs):
        nonlocal seen_prompt
        seen_prompt = kwargs["prompt"]
        return _make_engine_core_request(kwargs["request_id"])

    input_processor = mocker.Mock()
    input_processor.process_inputs.side_effect = process_inputs
    engine.input_processor = input_processor

    engine._build_add_request_message(
        request_id="req-no-replica",
        prompt={
            "prompt": "describe",
            "multi_modal_data": {"image": "same-image"},
        },
        sampling_params_list=[params],
        final_stage_id=0,
    )

    assert seen_prompt is not None
    assert "multi_modal_uuids" not in seen_prompt
    assert stage_pool.get_bound_replica_id("req-no-replica") is None


def test_stage_pool_replica_count_falls_back_to_clients():
    class PoolWithoutLiveNumReplicas:
        clients = [object(), None, object()]

    assert AsyncOmniEngine._stage_pool_replica_count(PoolWithoutLiveNumReplicas()) == 2


def test_stage_pool_is_distributed_falls_back_to_hub():
    class PoolWithoutIsDistributed:
        _hub = object()

    assert AsyncOmniEngine._stage_pool_is_distributed(PoolWithoutIsDistributed()) is True


def test_build_add_request_message_releases_preselected_replica_on_preprocess_error(mocker: MockerFixture):
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.model = "test-model"
    engine.default_sampling_params_list = [params]
    engine.stage_metadata = [StageRuntimeInfo(final_output=False, final_output_type=None, stage_type="llm")]
    engine.supported_tasks = ("generate",)
    stage_pool = StagePool(0, [_FakeStageClient(), _FakeStageClient()])
    engine.stage_pools = [stage_pool]

    input_processor = mocker.Mock()
    input_processor.process_inputs.side_effect = RuntimeError("boom")
    engine.input_processor = input_processor

    with pytest.raises(RuntimeError, match="boom"):
        engine._build_add_request_message(
            request_id="req-error",
            prompt={
                "prompt": "describe",
                "multi_modal_data": {"image": "same-image"},
            },
            sampling_params_list=[params],
            final_stage_id=0,
        )

    assert stage_pool.get_bound_replica_id("req-error") is None


def test_cfg_companion_build_failure_admits_nothing(mocker: MockerFixture):
    """A guided request is all-or-nothing.

    Admitting the parent before its companion is built leaves an orphan: a
    model whose guidance is mandatory never completes the pair, so the request
    holds scheduler and KV capacity for the scheduler's whole hold budget and
    then produces no audio. The build has to raise before anything is enqueued.
    """
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.prompt_expand_func = mocker.Mock(side_effect=ValueError("cannot build the null twin"))
    engine.supported_tasks = ("generate",)
    engine.input_processor = mocker.Mock()
    engine.request_queue = mocker.Mock()

    with pytest.raises(ValueError, match="null twin"):
        engine._build_cfg_companions(
            parent_id="req",
            original_prompt={"prompt": "positive"},
            stage0_params=params,
            sampling_params_list=[params],
        )

    engine.request_queue.sync_q.put.assert_not_called()


def test_cfg_companion_processing_failure_admits_nothing(mocker: MockerFixture):
    """The same holds when expansion succeeds but input processing does not."""
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.prompt_expand_func = lambda *_args: [
        ExpandedPrompt(
            prompt={"prompt": "negative"},
            role="cfg_text",
            request_id_suffix="__cfg_text",
        )
    ]
    engine.supported_tasks = ("generate",)
    engine.input_processor = mocker.Mock()
    engine.input_processor.process_inputs.side_effect = RuntimeError("tokenizer rejected the prompt")
    engine.request_queue = mocker.Mock()

    with pytest.raises(RuntimeError, match="tokenizer rejected"):
        engine._build_cfg_companions(
            parent_id="req",
            original_prompt={"prompt": "positive"},
            stage0_params=params,
            sampling_params_list=[params],
        )

    engine.request_queue.sync_q.put.assert_not_called()


def test_cfg_companion_build_returns_messages_without_enqueueing(mocker: MockerFixture):
    """Construction and admission are separate steps."""
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.prompt_expand_func = lambda *_args: [
        ExpandedPrompt(
            prompt={"prompt": "negative"},
            role="cfg_text",
            request_id_suffix="__cfg_text",
        )
    ]
    engine.supported_tasks = ("generate",)
    engine.input_processor = mocker.Mock()
    engine.input_processor.process_inputs.return_value = _make_engine_core_request("req__cfg_text")
    engine.request_queue = mocker.Mock()

    companions = engine._build_cfg_companions(
        parent_id="req",
        original_prompt={"prompt": "positive"},
        stage0_params=params,
        sampling_params_list=[params],
    )

    assert len(companions) == 1
    assert companions[0].companion_id == "req__cfg_text"
    assert companions[0].parent_id == "req"
    engine.request_queue.sync_q.put.assert_not_called()


def test_cfg_expansion_returning_nothing_is_not_an_error(mocker: MockerFixture):
    """An unguided request expands to no companions and admits normally."""
    engine = object.__new__(AsyncOmniEngine)
    params = SamplingParams(max_tokens=8)
    engine.prompt_expand_func = lambda *_args: []
    engine.supported_tasks = ("generate",)
    engine.input_processor = mocker.Mock()
    engine.request_queue = mocker.Mock()

    assert (
        engine._build_cfg_companions(
            parent_id="req",
            original_prompt={"prompt": "positive"},
            stage0_params=params,
            sampling_params_list=[params],
        )
        == []
    )
    engine.request_queue.sync_q.put.assert_not_called()