# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for HSDP (Hybrid Sharded Data Parallel) configuration and utilities."""

import gc
import os
import socket

import pytest
import torch
import torch.distributed as dist
import torch.nn as nn
from torch.distributed.tensor import DeviceMesh, DTensor

from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.distributed.hsdp import (
    HSDPInferenceConfig,
    _unshardable_parameters,
    shard_model,
)

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


class _PackedBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(2, 2))
        self.packed_weight = nn.Parameter(torch.ones(2, 2, dtype=torch.uint8), requires_grad=False)
        self.input_global_scale = nn.Parameter(torch.tensor(1.0), requires_grad=False)


class _PackedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.block = _PackedBlock()
        self.root_weight = nn.Parameter(torch.ones(2))


def _find_free_port() -> int:
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
        sock.bind(("127.0.0.1", 0))
        return int(sock.getsockname()[1])


@pytest.fixture(scope="module")
def cpu_process_group():
    if dist.is_initialized():
        yield
        return

    master_port = _find_free_port()
    os.environ.update(
        {
            "RANK": "0",
            "LOCAL_RANK": "0",
            "WORLD_SIZE": "1",
            "MASTER_ADDR": "127.0.0.1",
            "MASTER_PORT": str(master_port),
        }
    )
    dist.init_process_group("gloo", rank=0, world_size=1)
    try:
        yield
    finally:
        dist.destroy_process_group()
        for key in ("MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"):
            os.environ.pop(key, None)
        gc.collect()


def test_hsdp_keeps_packed_and_scalar_parameters_local(cpu_process_group):
    model = _PackedModel()
    ignored_params = _unshardable_parameters(model)
    expected_ignored = {model.block.packed_weight, model.block.input_global_scale}
    assert ignored_params == expected_ignored

    packed_weight = model.block.packed_weight
    input_global_scale = model.block.input_global_scale
    shard_model(
        model,
        mesh=DeviceMesh("cpu", [0]),
        hsdp_shard_conditions=[lambda name, _module: name == "block"],
        ignored_params=ignored_params,
    )

    assert isinstance(model.block.weight, DTensor)
    assert isinstance(model.root_weight, DTensor)
    assert model.block.packed_weight is packed_weight
    assert model.block.input_global_scale is input_global_scale
    assert not isinstance(model.block.packed_weight, DTensor)
    assert not isinstance(model.block.input_global_scale, DTensor)


class TestHSDPInferenceConfig:
    """Tests for HSDPInferenceConfig dataclass."""

    def test_default_values(self):
        """Test default configuration values."""
        config = HSDPInferenceConfig()
        assert config.enabled is False
        assert config.hsdp_replicate_size == 1
        assert config.hsdp_shard_size == -1

    def test_custom_values(self):
        """Test custom configuration values."""
        config = HSDPInferenceConfig(
            enabled=True,
            hsdp_replicate_size=2,
            hsdp_shard_size=4,
        )
        assert config.enabled is True
        assert config.hsdp_replicate_size == 2
        assert config.hsdp_shard_size == 4


class TestDiffusionParallelConfigHSDP:
    """Tests for HSDP settings in DiffusionParallelConfig."""

    def test_hsdp_disabled_by_default(self):
        """HSDP should be disabled by default."""
        config = DiffusionParallelConfig()
        assert config.use_hsdp is False
        assert config.hsdp_shard_size == -1
        assert config.hsdp_replicate_size == 1

    def test_hsdp_auto_shard_size(self):
        """Test auto-calculation of hsdp_shard_size when use_hsdp=True."""
        # ulysses_degree=4 -> world_size=4
        # hsdp_shard_size should be auto-calculated as 4 // 1 = 4
        config = DiffusionParallelConfig(
            ulysses_degree=4,
            use_hsdp=True,
        )
        assert config.world_size == 4
        assert config.hsdp_shard_size == 4
        assert config.hsdp_replicate_size == 1

    def test_hsdp_auto_shard_size_fails_standalone(self):
        """Test that auto-calculate fails when other parallelism is all 1."""
        # When all other parallelism is 1, cannot auto-calculate
        # User must specify hsdp_shard_size explicitly
        with pytest.raises(ValueError, match="Cannot auto-calculate hsdp_shard_size"):
            DiffusionParallelConfig(
                use_hsdp=True,
                # All other parallelism defaults to 1
            )

    def test_hsdp_standalone_mode(self):
        """Test standalone HSDP (HSDP without other parallelism)."""
        # Standalone HSDP: all other parallelism=1, explicit shard_size
        config = DiffusionParallelConfig(
            use_hsdp=True,
            hsdp_shard_size=4,  # Explicit shard size
            hsdp_replicate_size=1,
        )
        # world_size should be determined by HSDP
        assert config.world_size == 4
        assert config.hsdp_shard_size == 4
        assert config.hsdp_replicate_size == 1

    def test_hsdp_standalone_with_replicate(self):
        """Test standalone HSDP with replication."""
        config = DiffusionParallelConfig(
            use_hsdp=True,
            hsdp_shard_size=4,
            hsdp_replicate_size=2,
        )
        # world_size = shard_size * replicate_size
        assert config.world_size == 8
        assert config.hsdp_shard_size == 4
        assert config.hsdp_replicate_size == 2

    def test_hsdp_with_replicate(self):
        """Test HSDP with replication (hybrid mode) combined with other parallelism."""
        # world_size=8, replicate=2 -> shard_size should be 4
        config = DiffusionParallelConfig(
            ulysses_degree=8,
            use_hsdp=True,
            hsdp_replicate_size=2,
        )
        assert config.world_size == 8
        assert config.hsdp_shard_size == 4
        assert config.hsdp_replicate_size == 2

    def test_hsdp_explicit_shard_size_valid(self):
        """Test explicit hsdp_shard_size that matches world_size."""
        config = DiffusionParallelConfig(
            ulysses_degree=4,
            use_hsdp=True,
            hsdp_shard_size=4,
            hsdp_replicate_size=1,
        )
        assert config.hsdp_shard_size == 4

    def test_hsdp_explicit_shard_size_invalid(self):
        """Test that invalid HSDP dimensions raise an error when combined with other parallelism."""
        with pytest.raises(ValueError, match="HSDP dimensions"):
            DiffusionParallelConfig(
                ulysses_degree=4,  # world_size=4
                use_hsdp=True,
                hsdp_shard_size=3,  # 1 * 3 != 4
                hsdp_replicate_size=1,
            )

    def test_hsdp_replicate_size_exceeds_world_size(self):
        """Test that replicate_size > world_size raises an error."""
        with pytest.raises(ValueError, match="replicate_size.*must evenly divide world_size"):
            DiffusionParallelConfig(
                ulysses_degree=4,  # world_size=4
                use_hsdp=True,
                hsdp_replicate_size=8,  # 8 > 4, invalid
            )

    def test_hsdp_combined_world_size(self):
        """Test that combined HSDP matches other parallelism world_size."""
        config_no_hsdp = DiffusionParallelConfig(ulysses_degree=4)
        config_with_hsdp = DiffusionParallelConfig(
            ulysses_degree=4,
            use_hsdp=True,
            hsdp_shard_size=4,
        )
        # When combined with other parallelism, world_size should match
        assert config_no_hsdp.world_size == config_with_hsdp.world_size == 4

    def test_hsdp_standalone_world_size(self):
        """Test that standalone HSDP determines world_size."""
        config_hsdp = DiffusionParallelConfig(
            use_hsdp=True,
            hsdp_shard_size=8,
        )
        # Standalone HSDP: world_size is determined by HSDP
        assert config_hsdp.world_size == 8

    def test_hsdp_cannot_use_with_tp(self):
        """Test that HSDP and Tensor Parallelism cannot be used together."""
        with pytest.raises(ValueError, match="not compatible with TP"):
            DiffusionParallelConfig(
                tensor_parallel_size=2,
                use_hsdp=True,
                hsdp_shard_size=4,
            )

    def test_hsdp_cannot_use_with_dp(self):
        """Test that HSDP and Data Parallelism cannot be used together."""
        with pytest.raises(ValueError, match="not compatible with DP"):
            DiffusionParallelConfig(
                data_parallel_size=2,
                use_hsdp=True,
                hsdp_shard_size=4,
            )

    def test_from_dict_with_hsdp(self):
        """Test creating config from dict with HSDP settings."""
        config = DiffusionParallelConfig.from_dict(
            {
                "ulysses_degree": 4,
                "use_hsdp": True,
                "hsdp_replicate_size": 2,
            }
        )
        assert config.use_hsdp is True
        assert config.hsdp_replicate_size == 2
        assert config.hsdp_shard_size == 2  # auto: 4 // 2


class TestHSDPShardConditions:
    """Tests for _hsdp_shard_conditions matching logic."""

    @staticmethod
    def _is_transformer_block(name: str, module: nn.Module) -> bool:
        """Example shard condition matching transformer blocks."""
        return "blocks" in name and name.split(".")[-1].isdigit()

    def test_condition_matches_blocks(self):
        """Test that condition matches transformer block patterns."""
        cond = self._is_transformer_block
        # Should match
        assert cond("blocks.0", nn.Linear(10, 10)) is True
        assert cond("blocks.15", nn.Linear(10, 10)) is True
        assert cond("transformer.blocks.0", nn.Linear(10, 10)) is True
        # Should not match
        assert cond("blocks", nn.Linear(10, 10)) is False
        assert cond("blocks.norm", nn.Linear(10, 10)) is False
        assert cond("embeddings", nn.Linear(10, 10)) is False

    def test_model_with_shard_conditions(self):
        """Test model with _hsdp_shard_conditions attribute."""

        class MockModel(nn.Module):
            @staticmethod
            def _is_block(name: str, module: nn.Module) -> bool:
                return name.startswith("blocks.") and name.split(".")[-1].isdigit()

            _hsdp_shard_conditions = [_is_block]

            def __init__(self):
                super().__init__()
                self.blocks = nn.ModuleList([nn.Linear(10, 10) for _ in range(2)])

        model = MockModel()
        conditions = getattr(model, "_hsdp_shard_conditions", None)
        assert conditions is not None
        assert len(conditions) == 1

        # Verify conditions work on actual model modules
        matched = []
        for name, module in model.named_modules():
            if any(cond(name, module) for cond in conditions):
                matched.append(name)
        assert "blocks.0" in matched
        assert "blocks.1" in matched