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

"""
Unit tests for WorkerWrapperBase class.

This module tests the WorkerWrapperBase implementation:
- Initialization with and without worker extensions
- Custom pipeline initialization
- Method delegation via execute_method
- Attribute delegation via __getattr__
- Dynamic worker class extension
"""

from types import SimpleNamespace
from typing import Any

import pytest
from pytest_mock import MockerFixture

from vllm_omni.diffusion.worker.diffusion_worker import (
    CustomPipelineWorkerExtension,
    DiffusionWorker,
    WorkerWrapperBase,
)

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


# -------------------------------------------------------------------------
# Fixtures
# -------------------------------------------------------------------------


@pytest.fixture
def mock_od_config(mocker: MockerFixture):
    """Create a mock OmniDiffusionConfig for use in tests."""
    config = mocker.Mock()
    config.num_gpus = 1
    config.master_port = 12345
    config.enable_sleep_mode = False
    config.cache_backend = None
    config.cache_config = None
    config.model = "test-model"
    config.diffusion_load_format = "default"
    config.dtype = "float32"
    config.max_cpu_loras = 0
    config.lora_path = None
    config.lora_scale = 1.0
    return config


class TestExtension:
    """Simple test extension adding one custom method."""

    def custom_method(self):
        return "extension_method"


class MockCustomPipeline:
    """Mock custom pipeline for testing."""

    def __init__(self, *args, **kwargs):
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *args, **kwargs):
        return "pipeline_output"


# -------------------------------------------------------------------------
# Tests: Initialization
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseInitialization:
    """Test WorkerWrapperBase initialization behavior."""

    def test_basic_initialization(self, mocker: MockerFixture, mock_od_config):
        """Test basic initialization without extensions."""
        mock_worker_init = mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
        )

        assert wrapper.gpu_id == 0
        assert wrapper.od_config == mock_od_config
        assert wrapper.base_worker_class == DiffusionWorker
        assert wrapper.worker_extension_cls is None
        assert wrapper.custom_pipeline_args is None
        assert wrapper.worker is not None

        mock_worker_init.assert_called_once_with(
            local_rank=0,
            rank=0,
            od_config=mock_od_config,
            skip_load_model=False,
        )


# -------------------------------------------------------------------------
# Tests: Worker Extension Functionality
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseExtension:
    """Test WorkerWrapperBase worker extension functionality."""

    def test_prepare_worker_class_without_extension(self, mocker: MockerFixture, mock_od_config):
        """Test _prepare_worker_class without a worker extension."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
        )
        worker_class = wrapper._prepare_worker_class()
        assert worker_class == DiffusionWorker

    def test_prepare_worker_class_with_extension_class(self, mocker: MockerFixture, mock_od_config):
        """Test _prepare_worker_class with an explicit extension class."""

        class TestExtension:
            def custom_method(self):
                return "extension_method"

        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=TestExtension,
        )

        assert hasattr(wrapper.worker.__class__, "custom_method")
        assert TestExtension in wrapper.worker.__class__.__bases__

    def test_prepare_worker_class_with_extension_string(self, mocker: MockerFixture, mock_od_config):
        """Test _prepare_worker_class with worker extension as string."""
        mock_resolve = mocker.patch("vllm.utils.import_utils.resolve_obj_by_qualname")
        mock_resolve.return_value = TestExtension

        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls="tests.diffusion.test_worker_wrapper_base.TestExtension",
        )

        assert hasattr(wrapper.worker.__class__, "custom_method")


# -------------------------------------------------------------------------
# Tests: Method Delegation
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseDelegation:
    """Test WorkerWrapperBase delegation to wrapped worker."""

    def test_execute_model_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that execute_model() delegates to worker.execute_model()."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        mock_output = mocker.Mock()
        wrapper.worker.execute_model = mocker.Mock(return_value=mock_output)

        mock_req = mocker.Mock()
        result = wrapper.execute_model(mock_req, mock_od_config)

        wrapper.worker.execute_model.assert_called_once_with(mock_req, mock_od_config, kv_prefetch_job=None)
        assert result == mock_output

    def test_sleep_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that sleep() delegates to worker.sleep()."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.sleep = mocker.Mock(return_value=True)
        result = wrapper.sleep(level=1)

        wrapper.worker.sleep.assert_called_once_with(1)
        assert result is True

    def test_wake_up_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that wake_up() delegates to worker.wake_up()."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.wake_up = mocker.Mock(return_value=True)

        result = wrapper.wake_up(tags=["weights"])
        wrapper.worker.wake_up.assert_called_once_with(["weights"])
        assert result is True

    def test_shutdown_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test that shutdown() delegates to worker.shutdown()."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.shutdown = mocker.Mock(return_value=None)

        result = wrapper.shutdown()
        wrapper.worker.shutdown.assert_called_once()
        assert result is None

    def test_worker_shutdown_disables_offloader_before_distributed_teardown(self, mocker: MockerFixture):
        events: list[str] = []
        offload_backend = mocker.Mock()
        offload_backend.disable.side_effect = lambda: events.append("offload")
        kv_manager = mocker.Mock()
        kv_manager.shutdown_prefetch.side_effect = lambda: events.append("kv")
        destroy = mocker.patch(
            "vllm_omni.diffusion.worker.diffusion_worker.destroy_distributed_env",
            side_effect=lambda: events.append("distributed"),
        )
        worker = DiffusionWorker.__new__(DiffusionWorker)
        worker.model_runner = SimpleNamespace(
            offload_backend=offload_backend,
            kv_transfer_manager=kv_manager,
        )

        worker.shutdown()

        assert events == ["offload", "kv", "distributed"]
        destroy.assert_called_once_with()


# -------------------------------------------------------------------------
# Tests: execute_method
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseExecuteMethod:
    """Test WorkerWrapperBase.execute_method functionality."""

    def test_execute_method_success(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method successfully calls worker method."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.test_method = mocker.Mock(return_value="method_result")

        result = wrapper.execute_method("test_method", "arg1", kwarg1="value1")

        wrapper.worker.test_method.assert_called_once_with("arg1", kwarg1="value1")
        assert result == "method_result"

    def test_execute_method_with_no_args(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method with no arguments."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.no_args_method = mocker.Mock(return_value="no_args_result")

        result = wrapper.execute_method("no_args_method")
        wrapper.worker.no_args_method.assert_called_once_with()
        assert result == "no_args_result"

    def test_execute_method_error(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method raises exception on error."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.error_method = mocker.Mock(side_effect=RuntimeError("Test error"))

        with pytest.raises(RuntimeError, match="Test error"):
            wrapper.execute_method("error_method")

    def test_execute_method_invalid_type(self, mocker: MockerFixture, mock_od_config):
        """Test execute_method with invalid method type."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)

        with pytest.raises(AssertionError, match="Method must be str"):
            wrapper.execute_method(b"bytes_method")


# -------------------------------------------------------------------------
# Tests: __getattr__ delegation
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseGetAttr:
    """Test WorkerWrapperBase.__getattr__ delegation."""

    def test_getattr_delegation(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ delegates to worker attributes."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.custom_attribute = "test_value"
        assert wrapper.custom_attribute == "test_value"

    def test_getattr_method_access(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ delegates to worker methods."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        wrapper.worker.custom_method = mocker.Mock(return_value="method_result")

        result = wrapper.custom_method()
        wrapper.worker.custom_method.assert_called_once()
        assert result == "method_result"

    def test_getattr_missing_attribute(self, mocker: MockerFixture, mock_od_config):
        """Test __getattr__ raises AttributeError for missing attributes."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(gpu_id=0, od_config=mock_od_config, base_worker_class=DiffusionWorker)
        with pytest.raises(AttributeError):
            _ = wrapper.nonexistent_attribute


# -------------------------------------------------------------------------
# Tests: Edge Cases
# -------------------------------------------------------------------------


class TestWorkerWrapperBaseEdgeCases:
    """Test WorkerWrapperBase edge cases and special scenarios."""

    def test_extension_conflict_warning(self, mocker: MockerFixture, mock_od_config, caplog):
        """Test a warning is logged when an extension conflicts with worker."""
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        class ConflictExtension:
            def load_model(self):
                return "extension_load_model"

        mocker.patch.object(DiffusionWorker, "load_model")
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=ConflictExtension,
        )
        assert wrapper.worker is not None

    def test_multiple_extensions_same_class(self, mocker: MockerFixture, mock_od_config):
        """Test that applying same extension twice doesn't duplicate it."""

        class TestExtension:
            def custom_method(self):
                return "extension"

        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper1 = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=TestExtension,
        )
        wrapper2 = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=TestExtension,
        )

        assert hasattr(wrapper1.worker, "custom_method")
        assert hasattr(wrapper2.worker, "custom_method")


# -------------------------------------------------------------------------
# Tests: CustomPipelineWorkerExtension
# -------------------------------------------------------------------------


class TestCustomPipelineWorkerExtension:
    """Test CustomPipelineWorkerExtension functionality."""

    def test_re_init_pipeline_basic(self, mocker: MockerFixture, mock_od_config):
        """Test basic re_init_pipeline functionality."""
        mocker.patch("torch.accelerator.empty_cache")
        mocker.patch("gc.collect")
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=CustomPipelineWorkerExtension,
        )

        # Setup mock model_runner and pipeline
        mock_model_runner = mocker.Mock()
        mock_pipeline = mocker.Mock()
        mock_model_runner.pipeline = mock_pipeline
        wrapper.worker.model_runner = mock_model_runner
        wrapper.worker.init_lora_manager = mocker.Mock()
        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

        # Call re_init_pipeline
        wrapper.worker.re_init_pipeline(custom_args)

        # Verify load_model was called with correct arguments
        wrapper.worker.load_model.assert_called_once_with(
            load_format="custom_pipeline",
            custom_pipeline_name="tests.diffusion.test_worker_wrapper_base.MockCustomPipeline",
        )
        wrapper.worker.init_lora_manager.assert_called_once()

    def test_re_init_pipeline_cleanup(self, mocker: MockerFixture, mock_od_config):
        """Test that re_init_pipeline properly cleans up old pipeline."""
        mock_gc_collect = mocker.patch("gc.collect")
        mock_empty_cache = mocker.patch("torch.accelerator.empty_cache")
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=CustomPipelineWorkerExtension,
        )

        # Setup mock model_runner with pipeline
        mock_model_runner = mocker.Mock()
        mock_pipeline = mocker.Mock()
        mock_model_runner.pipeline = mock_pipeline
        wrapper.worker.model_runner = mock_model_runner
        wrapper.worker.init_lora_manager = mocker.Mock()
        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

        # Call re_init_pipeline
        wrapper.worker.re_init_pipeline(custom_args)

        # Verify cleanup was performed
        mock_gc_collect.assert_called_once()
        mock_empty_cache.assert_called_once()

    def test_re_init_pipeline_none_pipeline(self, mocker: MockerFixture, mock_od_config):
        """Test re_init_pipeline when pipeline is None."""
        mocker.patch("torch.accelerator.empty_cache")
        mocker.patch("gc.collect")
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=CustomPipelineWorkerExtension,
        )

        # Setup mock model_runner with None pipeline
        mock_model_runner = mocker.Mock()
        mock_model_runner.pipeline = None
        wrapper.worker.model_runner = mock_model_runner
        wrapper.worker.init_lora_manager = mocker.Mock()
        wrapper.worker.load_model = mocker.Mock()

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

        # Should not raise an error
        wrapper.worker.re_init_pipeline(custom_args)

        # Verify load_model was still called
        wrapper.worker.load_model.assert_called_once()

    def test_custom_pipeline_args_initialization(self, mocker: MockerFixture, mock_od_config):
        """Test initialization with custom_pipeline_args calls re_init_pipeline."""
        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        mock_prepare = mocker.patch.object(WorkerWrapperBase, "_prepare_worker_class")
        # Create a mock worker class with re_init_pipeline
        mock_worker_class = mocker.Mock()
        mock_worker_instance = mocker.Mock()
        mock_worker_instance.re_init_pipeline = mocker.Mock()
        mock_worker_class.return_value = mock_worker_instance
        mock_prepare.return_value = mock_worker_class

        _ = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            custom_pipeline_args=custom_args,
        )

        # Verify re_init_pipeline was called with custom_pipeline_args
        mock_worker_instance.re_init_pipeline.assert_called_once_with(custom_args)

    def test_custom_pipeline_with_explicit_extension(self, mocker: MockerFixture, mock_od_config):
        """Test that explicit worker_extension_cls is preserved when custom_pipeline_args is provided."""

        class CustomExtension:
            def re_init_pipeline(self, custom_pipeline_args: dict[str, Any]):
                return "custom_re_init_pipeline"

            def custom_extension_method(self):
                return "custom_extension_method"

        custom_args = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}

        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)
        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=CustomExtension,
            custom_pipeline_args=custom_args,
        )

        # Should still have the explicitly provided extension
        assert CustomExtension in wrapper.worker.__class__.__bases__
        assert hasattr(wrapper.worker, "custom_extension_method")

    def test_re_init_pipeline_multiple_calls(self, mocker: MockerFixture, mock_od_config):
        """Test calling re_init_pipeline multiple times."""
        mocker.patch("torch.accelerator.empty_cache")
        mocker.patch("gc.collect")
        mocker.patch.object(DiffusionWorker, "__init__", return_value=None)

        wrapper = WorkerWrapperBase(
            gpu_id=0,
            od_config=mock_od_config,
            base_worker_class=DiffusionWorker,
            worker_extension_cls=CustomPipelineWorkerExtension,
        )

        # Setup mock model_runner
        mock_model_runner = mocker.Mock()
        mock_pipeline1 = mocker.Mock()
        mock_pipeline2 = mocker.Mock()
        mock_model_runner.pipeline = mock_pipeline1
        wrapper.worker.model_runner = mock_model_runner
        wrapper.worker.init_lora_manager = mocker.Mock()
        wrapper.worker.load_model = mocker.Mock()

        # First call
        custom_args1 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
        wrapper.worker.re_init_pipeline(custom_args1)

        # Update pipeline for second call
        mock_model_runner.pipeline = mock_pipeline2

        # Second call
        custom_args2 = {"pipeline_class": "tests.diffusion.test_worker_wrapper_base.MockCustomPipeline"}
        wrapper.worker.re_init_pipeline(custom_args2)

        # Verify load_model was called twice with different pipelines
        assert wrapper.worker.load_model.call_count == 2
        assert wrapper.worker.init_lora_manager.call_count == 2