import unittest

import torch

from sglang.srt.distributed import init_distributed_environment
from sglang.srt.distributed.parallel_state import (
    destroy_distributed_environment,
    destroy_model_parallel,
    get_tp_group,
    initialize_model_parallel,
)
from sglang.srt.layers.dp_attention import set_dp_buffer_len
from sglang.srt.layers.moe.token_dispatcher.flashinfer import FlashinferDispatcher
from sglang.srt.layers.moe.utils import initialize_moe_config
from sglang.srt.runtime_context import get_context, publish
from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
from sglang.test.test_utils import CustomTestCase, publish_build_topology


class TestFlashinferDispatcher(CustomTestCase):
    @classmethod
    def setUpClass(cls):
        # Dist-init first: world_size (and so the tp/ep width ServerArgs must
        # carry) is only known after it, and init_distributed_environment
        # itself reads no published config.
        init_distributed_environment(
            world_size=-1,  # Auto-detect from environment
            rank=-1,  # Auto-detect from environment
            local_rank=-1,  # Auto-detect from environment
            backend="nccl",
        )
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        device = torch.device(f"cuda:{rank % torch.cuda.device_count()}")
        torch.cuda.set_device(device)

        server_args = ServerArgs(
            model_path="dummy", tp_size=world_size, ep_size=world_size
        )
        server_args.moe_runner_backend = "flashinfer_cutlass"
        server_args.moe_a2a_backend = "flashinfer"
        cls.server_args = server_args
        set_global_server_args_for_scheduler(server_args)
        publish(server_args, role="scheduler")
        initialize_moe_config()

        publish_build_topology(tp_size=world_size, ep_size=world_size, world_rank=rank)
        initialize_model_parallel()

    @classmethod
    def tearDownClass(cls):
        try:
            from flashinfer.comm.trtllm_moe_alltoall import MoeAlltoAll

            for workspace in MoeAlltoAll._WORKSPACE_CACHE.values():
                mnnvl_mem = workspace.get("mnnvl_mem")
                if mnnvl_mem is not None and "ptr" in vars(mnnvl_mem):
                    del mnnvl_mem.ptr
            MoeAlltoAll._WORKSPACE_CACHE.clear()
        except ImportError:
            pass
        destroy_model_parallel()
        destroy_distributed_environment()

    def create_dispatcher(
        self, router_topk=2, num_experts=8, num_local_experts=4, hidden_size=128
    ):
        """Helper to create dispatcher instance"""
        return FlashinferDispatcher(
            group=get_tp_group().device_group,
            router_topk=router_topk,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
            params_dtype=torch.bfloat16,
        )

    def set_dispatch_type(self, dispatch_type):
        get_context().override(
            "test_flashinfer_dispatcher",
            flashinfer_a2a_dispatch_type=dispatch_type,
        )

    def _zero_moe_a2a_dispatch_payloads(self):
        # Shared MoeAlltoAll workspaces keep stale recv payloads across tests.
        # Zero only the payload region so unused-source == 0 asserts stay valid.
        try:
            from flashinfer.comm.trtllm_moe_alltoall import (
                MoeAlltoAll,
                get_moe_alltoall_module,
            )
        except ImportError:
            return

        module = get_moe_alltoall_module()
        for ws in MoeAlltoAll._WORKSPACE_CACHE.values():
            workspace = ws["workspace"]
            aux = int(
                module.moe_a2a_get_aux_data_size(
                    ws["ep_size"],
                    ws["max_num_tokens"],
                    ws["eplb_stats_num_experts"],
                )
            )
            aux = ((aux + 127) // 128) * 128
            if aux < workspace.shape[1]:
                workspace[:, aux:].zero_()

    def test_dispatch_basic(self):
        """Test basic dispatch functionality"""
        self.set_dispatch_type("bf16")
        num_tokens = 16
        hidden_size = 128
        router_topk = 1  # Single expert per token for simplicity
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        num_experts = world_size
        num_local_experts = 1  # One expert per rank

        set_dp_buffer_len(
            global_dp_buffer_len=num_tokens * world_size,
            local_dp_buffer_len=num_tokens,
            dp_max_padding=True,
            global_num_tokens=None,
        )

        # Create tokens with rank number
        hidden_states = torch.full(
            (num_tokens, hidden_size), 100.0 + rank, dtype=torch.bfloat16, device="cuda"
        )

        # Route all tokens from rank i to expert (i+1) % world_size
        target_rank = (rank + 1) % world_size
        target_expert = target_rank  # Since we have 1 expert per rank

        topk_ids = torch.full(
            (num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
        )
        topk_weights = torch.ones(
            (num_tokens, router_topk), dtype=torch.float32, device="cuda"
        )

        from sglang.srt.layers.moe.topk import StandardTopKOutput

        topk_output = StandardTopKOutput(
            topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
        )

        torch.distributed.barrier()
        dispatcher = self.create_dispatcher(
            router_topk=router_topk,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
        )
        dispatcher.set_quant_config({"input_global_scale": None})

        dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
        received_hidden_states = dispatch_output.hidden_states
        self.assertEqual(dispatch_output.hidden_states_scale, None)

        # Expected: we should receive tokens from rank (rank - 1) % world_size
        expected_source_rank = (rank - 1 + world_size) % world_size

        # Verify we received the right number of tokens
        self.assertEqual(
            received_hidden_states.shape[0],
            num_tokens * world_size,
            f"Should receive {num_tokens * world_size} tokens",
        )

        # Verify tokens came from the expected source
        self.assertTrue(
            torch.all(
                received_hidden_states[
                    expected_source_rank * num_tokens : (expected_source_rank + 1)
                    * num_tokens
                ]
                == 100.0 + expected_source_rank
            )
        )
        self.assertTrue(
            torch.all(
                received_hidden_states[: expected_source_rank * num_tokens] == 0.0
            )
        )
        self.assertTrue(
            torch.all(
                received_hidden_states[(expected_source_rank + 1) * num_tokens :] == 0.0
            )
        )

    def test_dispatch_with_empty_tokens(self):
        """Test dispatch when there are no tokens (edge case)"""
        self.set_dispatch_type("bf16")
        # This tests the dummy token handling
        num_tokens = 16
        hidden_size = 128
        router_topk = 1  # Single expert per token for simplicity
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        num_experts = world_size
        num_local_experts = 1  # One expert per rank

        set_dp_buffer_len(
            global_dp_buffer_len=num_tokens * world_size,
            local_dp_buffer_len=num_tokens,
            dp_max_padding=False,
            global_num_tokens=[16, 0, 16, 16],
        )

        # Route all tokens from rank i to expert (i+1) % world_size
        target_rank = (rank + 1) % world_size
        target_expert = target_rank  # Since we have 1 expert per rank

        # Create tokens with rank number, rank 1 has no tokens
        if rank == 1:
            hidden_states = torch.empty(
                0, hidden_size, dtype=torch.bfloat16, device="cuda"
            )
            topk_ids = torch.empty(0, router_topk, dtype=torch.int32, device="cuda")
            topk_weights = torch.empty(
                0, router_topk, dtype=torch.float32, device="cuda"
            )
        else:
            hidden_states = torch.full(
                (num_tokens, hidden_size),
                100.0 + rank,
                dtype=torch.bfloat16,
                device="cuda",
            )
            topk_ids = torch.full(
                (num_tokens, router_topk),
                target_expert,
                dtype=torch.int32,
                device="cuda",
            )
            topk_weights = torch.ones(
                (num_tokens, router_topk), dtype=torch.float32, device="cuda"
            )

        from sglang.srt.layers.moe.topk import StandardTopKOutput

        topk_output = StandardTopKOutput(
            topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
        )

        self._zero_moe_a2a_dispatch_payloads()
        torch.distributed.barrier()

        dispatcher = self.create_dispatcher(
            router_topk=router_topk,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
        )
        dispatcher.set_quant_config({"input_global_scale": None})

        dispatch_output = dispatcher.dispatch(hidden_states, topk_output)
        received_hidden_states = dispatch_output.hidden_states

        # Expected: we should receive tokens from rank (rank - 1) % world_size
        expected_source_rank = (rank - 1 + world_size) % world_size

        # Verify we received the right number of tokens
        self.assertEqual(
            received_hidden_states.shape[0],
            num_tokens * world_size,
            f"Should receive {num_tokens * world_size} tokens",
        )

        # Verify tokens came from the expected source
        if rank == 2:
            # Rank 2 should receive no tokens since rank 1 was empty
            self.assertTrue(
                torch.all(received_hidden_states == 0.0),
                "Rank should receive no tokens",
            )
        else:
            self.assertTrue(
                torch.all(
                    received_hidden_states[
                        expected_source_rank * num_tokens : (expected_source_rank + 1)
                        * num_tokens
                    ]
                    == 100.0 + expected_source_rank
                ),
                "Rank {rank} should receive tokens from the expected source {expected_source_rank}",
            )
            self.assertTrue(
                torch.all(
                    received_hidden_states[: expected_source_rank * num_tokens] == 0.0
                ),
                "Rank should receive no tokens from previous ranks",
            )
            self.assertTrue(
                torch.all(
                    received_hidden_states[(expected_source_rank + 1) * num_tokens :]
                    == 0.0
                ),
                "Rank should receive no tokens from next ranks",
            )

    def test_dispatch_with_fp4_quantization(self):
        """Test dispatch with FP4 quantization enabled"""
        self.set_dispatch_type("nvfp4")
        num_tokens = 128
        hidden_size = 128
        router_topk = 1  # Single expert per token for simplicity
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        num_experts = world_size
        num_local_experts = 1  # One expert per rank

        set_dp_buffer_len(
            global_dp_buffer_len=num_tokens * world_size,
            local_dp_buffer_len=num_tokens,
            dp_max_padding=True,
            global_num_tokens=None,
        )

        # Create tokens with random values
        hidden_states = torch.randn(
            (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
        )

        # Route all tokens from rank i to expert (i+1) % world_size
        target_rank = (rank + 1) % world_size
        target_expert = target_rank  # Since we have 1 expert per rank

        topk_ids = torch.full(
            (num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
        )
        topk_weights = torch.ones(
            (num_tokens, router_topk), dtype=torch.float32, device="cuda"
        )

        from sglang.srt.layers.moe.topk import StandardTopKOutput

        topk_output = StandardTopKOutput(
            topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
        )

        dispatcher = self.create_dispatcher(
            router_topk=router_topk,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
        )
        # Set input global scale to enable FP4 quantization
        input_global_scale = torch.tensor(1.0, dtype=torch.float32, device="cuda")
        dispatcher.set_quant_config({"input_global_scale": input_global_scale})

        dispatch_output = dispatcher.dispatch(hidden_states, topk_output)

        self.assertEqual(
            dispatch_output.hidden_states.shape,
            (num_tokens * world_size, hidden_size // 2),
        )
        self.assertEqual(dispatch_output.hidden_states.dtype, torch.uint8)

        self.assertNotEqual(dispatch_output.hidden_states_scale, None)
        self.assertEqual(
            dispatch_output.hidden_states_scale.numel(),
            num_tokens * world_size * (hidden_size // 16),
        )
        self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)

    def test_dispatch_with_mxfp8_quantization(self):
        """Test dispatch with MXFP8 quantization enabled"""
        self.set_dispatch_type("mxfp8")
        num_tokens = 128
        hidden_size = 128
        router_topk = 1
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        num_experts = world_size
        num_local_experts = 1

        set_dp_buffer_len(
            global_dp_buffer_len=num_tokens * world_size,
            local_dp_buffer_len=num_tokens,
            dp_max_padding=True,
            global_num_tokens=None,
        )

        hidden_states = torch.randn(
            (num_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
        )

        target_rank = (rank + 1) % world_size
        target_expert = target_rank
        topk_ids = torch.full(
            (num_tokens, router_topk), target_expert, dtype=torch.int32, device="cuda"
        )
        topk_weights = torch.ones(
            (num_tokens, router_topk), dtype=torch.float32, device="cuda"
        )

        from sglang.srt.layers.moe.topk import StandardTopKOutput

        topk_output = StandardTopKOutput(
            topk_weights=topk_weights, topk_ids=topk_ids, router_logits=None
        )

        dispatcher = self.create_dispatcher(
            router_topk=router_topk,
            num_experts=num_experts,
            num_local_experts=num_local_experts,
            hidden_size=hidden_size,
        )
        dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})

        dispatch_output = dispatcher.dispatch(hidden_states, topk_output)

        self.assertEqual(
            dispatch_output.hidden_states.shape,
            (num_tokens * world_size, hidden_size),
        )
        self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
        self.assertEqual(dispatch_output.output_dtype, torch.bfloat16)

        self.assertIsNotNone(dispatch_output.hidden_states_scale)
        self.assertEqual(
            dispatch_output.hidden_states_scale.shape,
            (num_tokens * world_size, hidden_size // 32),
        )
        self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
        self.assertEqual(
            dispatch_output.topk_output.topk_ids.shape,
            (num_tokens * world_size, router_topk),
        )
        self.assertEqual(dispatch_output.topk_output.topk_ids.dtype, torch.int32)

    def test_dispatch_with_mxfp8_quantization_and_empty_rank(self):
        """All ranks must contribute the same payload dtypes, including empty ranks."""
        self.set_dispatch_type("mxfp8")
        num_tokens = 16
        hidden_size = 128
        router_topk = 1
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()
        empty_rank = 1

        global_num_tokens = [num_tokens] * world_size
        global_num_tokens[empty_rank] = 0
        set_dp_buffer_len(
            global_dp_buffer_len=num_tokens * world_size,
            local_dp_buffer_len=num_tokens,
            dp_max_padding=False,
            global_num_tokens=global_num_tokens,
        )

        local_tokens = 0 if rank == empty_rank else num_tokens
        hidden_states = torch.randn(
            (local_tokens, hidden_size), dtype=torch.bfloat16, device="cuda"
        )
        target_expert = (rank + 1) % world_size
        topk_ids = torch.full(
            (local_tokens, router_topk),
            target_expert,
            dtype=torch.int32,
            device="cuda",
        )
        topk_weights = torch.ones(
            (local_tokens, router_topk), dtype=torch.float32, device="cuda"
        )

        from sglang.srt.layers.moe.topk import StandardTopKOutput

        dispatcher = self.create_dispatcher(
            router_topk=router_topk,
            num_experts=world_size,
            num_local_experts=1,
            hidden_size=hidden_size,
        )
        dispatcher.set_quant_config({"input_global_scale": None, "use_mxfp8": True})
        self._zero_moe_a2a_dispatch_payloads()
        torch.distributed.barrier()
        dispatch_output = dispatcher.dispatch(
            hidden_states,
            StandardTopKOutput(
                topk_weights=topk_weights,
                topk_ids=topk_ids,
                router_logits=None,
            ),
        )

        self.assertEqual(dispatch_output.hidden_states.dtype, torch.float8_e4m3fn)
        self.assertEqual(dispatch_output.hidden_states_scale.dtype, torch.uint8)
        self.assertEqual(
            dispatch_output.hidden_states.shape,
            (num_tokens * world_size, hidden_size),
        )


if __name__ == "__main__":
    """
    Usage
    torchrun --nproc_per_node=4 test_flashinfer_dispatcher.py
    """
    unittest.main()