# 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="cannot be used with TP or DP"):
            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="cannot be used with TP or 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 TestStandaloneHSDPDetection:
    """Tests for standalone HSDP detection and dit_parallel_size calculation.

    These tests verify the logic used in initialize_model_parallel() to detect
    standalone HSDP mode and calculate effective parallel sizes.

    Standalone HSDP is when all non-HSDP parallelism dimensions are 1.
    """

    @staticmethod
    def compute_standalone_hsdp_params(
        data_parallel_size: int = 1,
        cfg_parallel_size: int = 1,
        sequence_parallel_size: int = 1,
        pipeline_parallel_size: int = 1,
        tensor_parallel_size: int = 1,
        fully_shard_degree: int = 1,
        hsdp_replicate_size: int = 1,
    ) -> dict:
        """Compute standalone HSDP detection parameters.

        This mirrors the logic in initialize_model_parallel().
        """
        dit_parallel_size = (
            data_parallel_size
            * cfg_parallel_size
            * sequence_parallel_size
            * pipeline_parallel_size
            * tensor_parallel_size
        )

        # Check for standalone HSDP: all non-HSDP parallelism dimensions are 1
        is_standalone_hsdp = dit_parallel_size == 1 and fully_shard_degree > 1

        # For standalone HSDP: use (fully_shard_degree * hsdp_replicate_size)
        if is_standalone_hsdp:
            effective_dit_parallel_size = fully_shard_degree * hsdp_replicate_size
        else:
            effective_dit_parallel_size = dit_parallel_size

        effective_dp_size = (fully_shard_degree * hsdp_replicate_size) if is_standalone_hsdp else data_parallel_size

        return {
            "original_dit_parallel_size": dit_parallel_size,
            "is_standalone_hsdp": is_standalone_hsdp,
            "effective_dit_parallel_size": effective_dit_parallel_size,
            "effective_dp_size": effective_dp_size,
        }

    def test_standalone_hsdp_basic(self):
        """Test basic standalone HSDP detection (shard_size=4, replicate=1)."""
        result = self.compute_standalone_hsdp_params(
            fully_shard_degree=4,
            hsdp_replicate_size=1,
        )
        assert result["original_dit_parallel_size"] == 1
        assert result["is_standalone_hsdp"] is True
        assert result["effective_dit_parallel_size"] == 4
        assert result["effective_dp_size"] == 4

    def test_standalone_hsdp_with_replicate(self):
        """Test standalone HSDP with replication (shard_size=4, replicate=2)."""
        result = self.compute_standalone_hsdp_params(
            fully_shard_degree=4,
            hsdp_replicate_size=2,
        )
        assert result["original_dit_parallel_size"] == 1
        assert result["is_standalone_hsdp"] is True
        assert result["effective_dit_parallel_size"] == 8  # 4 * 2
        assert result["effective_dp_size"] == 8

    def test_combined_hsdp_sp_not_standalone(self):
        """Test HSDP combined with SP is NOT detected as standalone.

        This is a regression test for the bug where the condition
        `dit_parallel_size == fully_shard_degree` incorrectly matched
        combined modes like SP=4 + HSDP=4.
        """
        result = self.compute_standalone_hsdp_params(
            sequence_parallel_size=4,
            fully_shard_degree=4,
            hsdp_replicate_size=1,
        )
        assert result["original_dit_parallel_size"] == 4
        assert result["is_standalone_hsdp"] is False
        # Should NOT override dp_size for combined mode
        assert result["effective_dp_size"] == 1  # original data_parallel_size

    def test_combined_hsdp_cfg_not_standalone(self):
        """Test HSDP combined with CFG is NOT detected as standalone."""
        result = self.compute_standalone_hsdp_params(
            cfg_parallel_size=2,
            fully_shard_degree=4,
            hsdp_replicate_size=1,
        )
        assert result["original_dit_parallel_size"] == 2
        assert result["is_standalone_hsdp"] is False
        assert result["effective_dp_size"] == 1

    def test_combined_hsdp_pp_not_standalone(self):
        """Test HSDP combined with PP is NOT detected as standalone."""
        result = self.compute_standalone_hsdp_params(
            pipeline_parallel_size=2,
            fully_shard_degree=4,
            hsdp_replicate_size=1,
        )
        assert result["original_dit_parallel_size"] == 2
        assert result["is_standalone_hsdp"] is False
        assert result["effective_dp_size"] == 1

    def test_no_hsdp_not_standalone(self):
        """Test that no HSDP (fully_shard_degree=1) is NOT standalone."""
        result = self.compute_standalone_hsdp_params(
            fully_shard_degree=1,
        )
        assert result["original_dit_parallel_size"] == 1
        assert result["is_standalone_hsdp"] is False
        assert result["effective_dp_size"] == 1

    def test_combined_multiple_parallelism_not_standalone(self):
        """Test HSDP combined with multiple parallelism is NOT standalone."""
        result = self.compute_standalone_hsdp_params(
            sequence_parallel_size=2,
            cfg_parallel_size=2,
            fully_shard_degree=4,
            hsdp_replicate_size=1,
        )
        assert result["original_dit_parallel_size"] == 4  # 2 * 2
        assert result["is_standalone_hsdp"] is False
        assert result["effective_dp_size"] == 1

    def test_standalone_hsdp_large_shard(self):
        """Test standalone HSDP with large shard size."""
        result = self.compute_standalone_hsdp_params(
            fully_shard_degree=8,
            hsdp_replicate_size=1,
        )
        assert result["is_standalone_hsdp"] is True
        assert result["effective_dit_parallel_size"] == 8
        assert result["effective_dp_size"] == 8

    def test_standalone_hsdp_large_replicate(self):
        """Test standalone HSDP with large replicate size."""
        result = self.compute_standalone_hsdp_params(
            fully_shard_degree=4,
            hsdp_replicate_size=4,
        )
        assert result["is_standalone_hsdp"] is True
        assert result["effective_dit_parallel_size"] == 16  # 4 * 4
        assert result["effective_dp_size"] == 16


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