"""Unit tests for vllm_omni.entrypoints.utils module."""

import os
from collections import Counter
from dataclasses import dataclass

import pytest
import torch
from pytest_mock import MockerFixture
from vllm.sampling_params import RequestOutputKind, SamplingParams

from vllm_omni.config.yaml_util import create_config
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.entrypoints.utils import (
    _convert_dataclasses_to_dict,
    _filter_dict_like_object,
    coerce_param_message_types,
    filter_dataclass_kwargs,
    filter_stages,
    load_and_resolve_stage_configs,
    load_stage_configs_from_yaml,
    resolve_model_config_path,
)

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


class TestFilterDictLikeObject:
    """Test suite for _filter_dict_like_object function."""

    def test_simple_dict(self):
        """Test filtering a simple dictionary with no callables."""
        input_dict = {"key1": "value1", "key2": 42, "key3": [1, 2, 3]}
        result = _filter_dict_like_object(input_dict)

        assert result == input_dict
        assert isinstance(result, dict)

    def test_dict_with_nested_values(self):
        """Test filtering dict with nested dict and list values."""
        input_dict = {
            "level1": {
                "level2": {"key": "value"},
                "list": [1, 2, 3],
            },
            "simple": "string",
        }

        result = _filter_dict_like_object(input_dict)

        # Nested dicts and lists should be recursively processed
        assert result["simple"] == "string"
        assert isinstance(result["level1"], dict)

    def test_dict_with_dataclass_values(self):
        """Test filtering dict containing dataclass values."""

        @dataclass
        class TestDataclass:
            field1: str
            field2: int

        obj = TestDataclass(field1="test", field2=42)
        input_dict = {"data": obj, "normal": "value"}

        result = _filter_dict_like_object(input_dict)

        # Dataclass should be converted to dict by recursive _convert_dataclasses_to_dict
        assert "data" in result
        assert "normal" in result
        assert result["normal"] == "value"

    def test_dict_with_counter_values(self):
        """Test filtering dict containing Counter objects."""
        counter_obj = Counter({"a": 1, "b": 2})
        input_dict = {"counter": counter_obj, "normal": "value"}

        result = _filter_dict_like_object(input_dict)

        # Counter should be converted to regular dict
        assert "counter" in result
        assert "normal" in result
        assert result["normal"] == "value"

    def test_empty_dict(self):
        """Test filtering an empty dictionary."""
        result = _filter_dict_like_object({})
        assert result == {}
        assert isinstance(result, dict)

    def test_dict_with_set_values(self):
        """Test filtering dict with set values."""
        input_dict = {"set_key": {1, 2, 3}, "normal": "value"}

        result = _filter_dict_like_object(input_dict)

        assert "set_key" in result
        assert "normal" in result
        # Set should be converted to list by _convert_dataclasses_to_dict
        assert result["normal"] == "value"

    def test_dict_with_none_values(self):
        """Test filtering dict with None values."""
        input_dict = {"key1": None, "key2": "value", "key3": 0}

        result = _filter_dict_like_object(input_dict)

        assert result == input_dict

    def test_dict_with_mixed_types(self):
        """Test filtering dict with mixed value types."""
        input_dict = {
            "string": "hello",
            "int": 42,
            "float": 3.14,
            "bool": True,
            "none": None,
            "list": [1, 2, 3],
            "tuple": (1, 2, 3),
            "set": {1, 2, 3},
            "dict": {"nested": "value"},
        }

        result = _filter_dict_like_object(input_dict)

        assert "string" in result
        assert "int" in result
        assert "float" in result
        assert "bool" in result
        assert "none" in result
        assert "list" in result
        assert "tuple" in result
        assert "set" in result
        assert "dict" in result

    def test_dict_preserves_key_types(self):
        """Test that dict key types are preserved."""
        input_dict = {
            "string_key": "value1",
            42: "value2",
            (1, 2): "value3",  # tuple as key
        }

        result = _filter_dict_like_object(input_dict)

        # Keys should remain the same
        assert "string_key" in result
        assert 42 in result
        assert (1, 2) in result

    def test_dict_with_recursive_structure(self, mocker: MockerFixture):
        """Test filtering dict with recursive/complex nested structure."""
        input_dict = {
            "level1": {
                "level2": {
                    "level3": {"key": "value"},
                    "callable": lambda x: x,
                }
            },
            "normal": "value",
        }

        mocker.patch("vllm_omni.entrypoints.utils.logger")
        result = _filter_dict_like_object(input_dict)

        # Normal key should exist
        assert "normal" in result
        # Level1 should exist
        assert "level1" in result

    def test_integration_with_convert_dataclasses(self, mocker: MockerFixture):
        """Test that _filter_dict_like_object integrates properly with _convert_dataclasses_to_dict."""

        @dataclass
        class Config:
            name: str
            count: int

        input_dict = {
            "config": Config(name="test", count=5),
            "func": lambda x: x,
            "normal": "value",
        }

        mocker.patch("vllm_omni.entrypoints.utils.logger")
        result = _filter_dict_like_object(input_dict)

        # Callable should be filtered
        assert "func" not in result
        # Config should be converted to dict
        assert "config" in result
        assert "normal" in result


class TestConvertDataclassesToDict:
    """Test suite for _convert_dataclasses_to_dict function."""

    def test_uses_filter_dict_like_object(self, mocker: MockerFixture):
        """Test that _convert_dataclasses_to_dict uses _filter_dict_like_object for dicts."""
        input_dict = {
            "normal": "value",
            "callable": lambda x: x,
        }

        mocker.patch("vllm_omni.entrypoints.utils.logger")
        result = _convert_dataclasses_to_dict(input_dict)

        # Callable should be filtered out by _filter_dict_like_object
        assert "normal" in result
        assert "callable" not in result


class TestFilterDataclassKwargs:
    """Test basic functionality of filter_dataclass_kwargs."""

    def test_simple_filtering(self):
        """Test basic dataclass kwargs filtering."""

        @dataclass
        class SimpleConfig:
            name: str
            count: int

        kwargs = {"name": "test", "count": 42, "invalid": "should_be_removed"}
        result = filter_dataclass_kwargs(SimpleConfig, kwargs)

        assert "name" in result
        assert "count" in result
        assert "invalid" not in result

    def test_invalid_dataclass_raises_error(self):
        """Test that non-dataclass raises ValueError."""
        with pytest.raises(ValueError, match="is not a dataclass"):
            filter_dataclass_kwargs(dict, {})

    def test_invalid_kwargs_type_raises_error(self):
        """Test that non-dict kwargs raises ValueError."""

        @dataclass
        class SimpleConfig:
            name: str

        with pytest.raises(ValueError, match="kwargs must be a dictionary"):
            filter_dataclass_kwargs(SimpleConfig, "invalid")

    def test_filters_omni_engine_args_unknown_fields(self):
        """Test that OmniEngineArgs kwargs are filtered to valid fields only."""
        kwargs = {
            "model": "dummy",
            "stage_id": 1,
            "engine_output_type": "image",
            "unknown_field": "drop_me",
        }

        result = filter_dataclass_kwargs(OmniEngineArgs, kwargs)

        assert "model" in result
        assert "stage_id" in result
        assert "engine_output_type" in result
        assert "unknown_field" not in result

    def test_filters_omni_diffusion_config_union_dataclass(self):
        """Test that OmniDiffusionConfig filters nested dataclass in Union fields."""
        kwargs = {
            "model": "dummy",
            "cache_config": {
                "rel_l1_thresh": 0.3,
                "extra_param": "should_drop",
            },
            "unknown_top": "drop_me",
        }

        result = filter_dataclass_kwargs(OmniDiffusionConfig, kwargs)

        assert "model" in result
        assert "cache_config" in result
        assert "unknown_top" not in result
        assert result["cache_config"]["rel_l1_thresh"] == 0.3
        assert "extra_param" not in result["cache_config"]


class TestResolveModelConfigPath:
    """Test suite for resolve_model_config_path function with diffusers format models."""

    def test_glm_image_diffusers_format_resolution(self, mocker: MockerFixture):
        """Test GlmImagePipeline diffusers class resolves to glm_image config."""
        mocker.patch(
            "vllm_omni.entrypoints.utils.file_or_path_exists",
            return_value=True,
        )
        mocker.patch(
            "vllm_omni.entrypoints.utils._try_get_class_name_from_diffusers_config",
            return_value="GlmImagePipeline",
        )
        mocker.patch(
            "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path",
            return_value="vllm_omni/model_executor/stage_configs",
        )

        original_exists = os.path.exists

        def mock_exists(path):
            if "glm_image.yaml" in str(path):
                return True
            return original_exists(path)

        mocker.patch("os.path.exists", side_effect=mock_exists)

        result = resolve_model_config_path("zai-org/GLM-Image")

        assert result is not None
        assert "glm_image.yaml" in result


class TestLoadAndResolveStageConfigs:
    def test_load_and_resolve_with_kwargs(self, mocker: MockerFixture):
        """Ensure that dtype survives default stage creation."""
        kwargs = {"dtype": torch.float32}
        mocker.patch("vllm_omni.entrypoints.utils.resolve_model_config_path", return_value=None)
        mocker.patch("vllm_omni.entrypoints.utils.load_stage_configs_from_model", return_value=([], None))

        config_path, stage_configs, _ = load_and_resolve_stage_configs(
            model="black-forest-labs/FLUX.2-klein-4B",
            stage_configs_path=None,
            kwargs=kwargs,
            trust_remote_code=False,
            default_stage_cfg_factory=lambda: AsyncOmniEngine._create_default_diffusion_stage_cfg(kwargs),
        )
        assert config_path is None
        assert len(stage_configs) == 1
        assert "dtype" in stage_configs[0]["engine_args"]

    def test_stage_configs_path_promotes_new_deploy_yaml_without_expanding_replicas(
        self, tmp_path, mocker: MockerFixture
    ):
        deploy_path = tmp_path / "qwen3_multi.yaml"
        deploy_path.write_text(
            'stages:\n  - stage_id: 0\n    devices: "0"\n  - stage_id: 1\n    devices: "1,2,3"\n    num_replicas: 3\n',
            encoding="utf-8",
        )

        returned_stage_configs = [
            create_config({"stage_id": 0, "runtime": {"devices": "0"}, "engine_args": {"model": "dummy"}}),
            create_config(
                {
                    "stage_id": 1,
                    "runtime": {"devices": "1,2,3", "num_replicas": 3},
                    "engine_args": {"model": "dummy"},
                }
            ),
        ]
        load_stage_configs = mocker.patch(
            "vllm_omni.entrypoints.utils.load_stage_configs_from_model",
            return_value=(returned_stage_configs, None),
        )

        config_path, stage_configs, _ = load_and_resolve_stage_configs(
            model="dummy-model",
            stage_configs_path=str(deploy_path),
            kwargs={},
            trust_remote_code=True,
        )

        load_stage_configs.assert_called_once_with(
            "dummy-model",
            trust_remote_code=True,
            base_engine_args={},
            deploy_config_path=str(deploy_path),
            stage_overrides=None,
            strategy_config_path=None,
        )
        assert config_path == str(deploy_path)
        assert len(stage_configs) == 2
        assert stage_configs[1].runtime.num_replicas == 3
        assert stage_configs[1].runtime.devices == "1,2,3"

    def test_filter_stages_selects_mode_stages_without_mutating_stage_config(self, tmp_path):
        config_path = tmp_path / "deploy.yaml"
        config_path.write_text(
            """modes:
  - mode: text-to-text
    stages: [0]
  - mode: text-to-image
    stages: [0, 1]
""",
            encoding="utf-8",
        )
        stages = [
            create_config(
                {
                    "stage_id": 0,
                    "runtime": {"requires_multimodal_data": True},
                    "final_output": False,
                    "final_output_type": None,
                }
            ),
            create_config(
                {
                    "stage_id": 1,
                    "runtime": {"requires_multimodal_data": True},
                    "final_output": True,
                    "final_output_type": "image",
                }
            ),
        ]

        filtered = filter_stages(str(config_path), stages, {"mode": "text-to-text"})

        assert len(filtered) == 1
        assert filtered[0].stage_id == 0
        assert filtered[0].runtime.requires_multimodal_data is True
        assert filtered[0].final_output is False
        assert filtered[0].final_output_type is None


class TestLoadStageConfigsFromYaml:
    """Regression tests for stage-config loading and merging."""

    def test_deep_merges_stage_engine_args(self, mocker: MockerFixture):
        yaml_config = create_config(
            {
                "async_chunk": True,
                "stage_args": [
                    {
                        "stage_id": 0,
                        "runtime": {"device": 0},
                        "engine_args": {
                            "parallel_config": {"tensor_parallel_size": 4},
                        },
                    }
                ],
            }
        )
        mocker.patch(
            "vllm_omni.entrypoints.utils.load_yaml_config",
            return_value=yaml_config,
        )

        stages = load_stage_configs_from_yaml(
            "fake.yaml",
            base_engine_args={
                "parallel_config": {
                    "tensor_parallel_size": 1,
                    "pipeline_parallel_size": 2,
                },
                "model": "base-model",
            },
        )

        merged_engine_args = stages[0]["engine_args"]
        assert merged_engine_args["parallel_config"]["tensor_parallel_size"] == 4
        assert merged_engine_args["parallel_config"]["pipeline_parallel_size"] == 2
        assert merged_engine_args["model"] == "base-model"
        assert merged_engine_args["async_chunk"] is True

    def test_merges_nested_stage_engine_args(self, mocker: MockerFixture):
        yaml_config = create_config(
            {
                "stage_args": [
                    {
                        "stage_id": 0,
                        "engine_args": {
                            "nested": {"override": 2},
                        },
                    }
                ],
            }
        )
        mocker.patch(
            "vllm_omni.entrypoints.utils.load_yaml_config",
            return_value=yaml_config,
        )

        stages = load_stage_configs_from_yaml(
            "fake.yaml",
            base_engine_args={"nested": {"base": 1}},
        )

        assert stages[0]["engine_args"]["nested"]["base"] == 1
        assert stages[0]["engine_args"]["nested"]["override"] == 2


class TestCumulativeStreamingCoercion:
    @pytest.mark.parametrize("skip_clone", [True, False])
    def test_cumulative_default_becomes_delta_if_stream(self, skip_clone):
        """Ensure cumulative messages are coercible to delta if streaming."""
        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
        sp.skip_clone = skip_clone
        result = coerce_param_message_types([sp], is_streaming=True)[0]
        assert isinstance(result, SamplingParams)
        assert result.output_kind == RequestOutputKind.DELTA
        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)

    @pytest.mark.parametrize("skip_clone", [True, False])
    def test_cumulative_default_becomes_final_only_if_not_stream(self, skip_clone):
        """Ensure cumulative messages are coercible to final only if not streaming."""
        sp = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
        sp.skip_clone = skip_clone
        result = coerce_param_message_types([sp], is_streaming=False)[0]
        assert isinstance(result, SamplingParams)
        assert result.output_kind == RequestOutputKind.FINAL_ONLY
        assert (skip_clone and sp is result) or (not skip_clone and sp is not result)

    @pytest.mark.parametrize("is_streaming", [True, False])
    @pytest.mark.parametrize("output_kind", [RequestOutputKind.DELTA, RequestOutputKind.FINAL_ONLY])
    def test_non_cumulative_are_coerced(self, output_kind, is_streaming):
        """Ensure non-cumulative params are coerced to the target type."""
        sp = SamplingParams(output_kind=output_kind)
        expected = RequestOutputKind.DELTA if is_streaming else RequestOutputKind.FINAL_ONLY
        result = coerce_param_message_types([sp], is_streaming=is_streaming)[0]
        assert isinstance(result, SamplingParams)
        assert result.output_kind == expected

    def test_coercion_applies_to_all_stages(self):
        """Ensure all stages are coerced to DELTA for streaming."""
        sp0 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
        sp1 = SamplingParams(output_kind=RequestOutputKind.CUMULATIVE)
        result = coerce_param_message_types([sp0, sp1], is_streaming=True)
        assert all([isinstance(r, SamplingParams) for r in result])
        assert result[0].output_kind == RequestOutputKind.DELTA
        assert result[1].output_kind == RequestOutputKind.DELTA