import argparse
import inspect
import os
import pathlib
import random
import sys
from typing import List, Tuple

import torch
import torch.distributed as dist


REPO_ROOT = pathlib.Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

import deep_gemm
from deep_gemm.utils.math import ceil_div


def parse_tokens_list(value: str) -> List[int]:
    return [int(item) for item in value.split(',') if item]


def gather_same_shape(tensor: torch.Tensor, group: dist.ProcessGroup) -> torch.Tensor:
    gathered = [torch.empty_like(tensor) for _ in range(dist.get_world_size(group))]
    dist.all_gather(gathered, tensor.contiguous(), group=group)
    return torch.stack(gathered, dim=0)


def init_test_dist(local_rank_arg: int = None) -> Tuple[int, int, dist.ProcessGroup]:
    local_rank = local_rank_arg if local_rank_arg is not None else int(os.environ.get('LOCAL_RANK', '0'))
    rank = int(os.environ.get('RANK', '0'))
    world_size = int(os.environ.get('WORLD_SIZE', '1'))
    master_addr = os.environ.get('MASTER_ADDR', '127.0.0.1')
    master_port = int(os.environ.get('MASTER_PORT', '8361'))

    torch.cuda.set_device(local_rank)
    sig = inspect.signature(dist.init_process_group)
    params = {
        'backend': 'nccl',
        'init_method': f'tcp://{master_addr}:{master_port}',
        'world_size': world_size,
        'rank': rank,
    }
    if 'device_id' in sig.parameters:
        params['device_id'] = torch.device(f'cuda:{local_rank}')
    dist.init_process_group(**params)
    torch.set_default_device('cuda')
    return rank, world_size, dist.new_group(list(range(world_size)))


def get_block_m(num_tokens: int, num_ranks: int, num_topk: int, num_experts: int) -> int:
    expected = num_tokens * num_ranks * num_topk / num_experts
    if expected <= 16.5:
        return 32
    if expected <= 64.5:
        return 64
    return 128


def make_topk(num_tokens: int, num_experts: int, num_topk: int, rank_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
    token_idx = torch.arange(num_tokens, device='cuda', dtype=torch.long).unsqueeze(1)
    topk_slot = torch.arange(num_topk, device='cuda', dtype=torch.long).unsqueeze(0)
    topk_idx = (token_idx * num_topk + topk_slot + rank_idx * 3) % num_experts
    topk_weights = (token_idx.float() * 0.125 + topk_slot.float() * 0.25 + rank_idx + 1.0).contiguous()
    return topk_idx.contiguous(), topk_weights.contiguous()


def make_weights(num_experts_per_rank: int, hidden: int, intermediate_hidden: int):
    l1_weights = torch.randn(
        (num_experts_per_rank, intermediate_hidden * 2, hidden),
        dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
    l2_weights = torch.randn(
        (num_experts_per_rank, hidden, intermediate_hidden),
        dtype=torch.float32, device='cuda').to(torch.float8_e4m3fn)
    l1_weights_sf = torch.ones(
        (num_experts_per_rank, ceil_div(intermediate_hidden * 2, 128), hidden // 128),
        dtype=torch.float32, device='cuda')
    l2_weights_sf = torch.ones(
        (num_experts_per_rank, ceil_div(hidden, 128), intermediate_hidden // 128),
        dtype=torch.float32, device='cuda')
    return deep_gemm.transform_weights_for_mega_moe(
        (l1_weights, l1_weights_sf), (l2_weights, l2_weights_sf))


def build_expected_entries(all_topk_idx: torch.Tensor,
                           rank_idx: int,
                           num_experts_per_rank: int) -> Tuple[List[List[Tuple[int, int, int]]], torch.Tensor]:
    num_ranks, num_tokens, num_topk = all_topk_idx.shape
    local_start = rank_idx * num_experts_per_rank
    local_end = local_start + num_experts_per_rank
    entries: List[List[Tuple[int, int, int]]] = [[] for _ in range(num_experts_per_rank)]
    counts = [0 for _ in range(num_experts_per_rank)]

    all_topk_idx_cpu = all_topk_idx.cpu()
    for src_rank in range(num_ranks):
        for token_idx in range(num_tokens):
            for topk_idx in range(num_topk):
                expert_idx = int(all_topk_idx_cpu[src_rank, token_idx, topk_idx])
                if local_start <= expert_idx < local_end:
                    local_expert = expert_idx - local_start
                    entries[local_expert].append((src_rank, token_idx, topk_idx))
                    counts[local_expert] += 1
    return entries, torch.tensor(counts, dtype=torch.int64, device='cuda')


def verify_case(buffer: deep_gemm.SymmBuffer,
                all_x_f32: torch.Tensor,
                all_x_sf: torch.Tensor,
                all_topk_idx: torch.Tensor,
                all_topk_weights: torch.Tensor,
                cumulative_stats: torch.Tensor,
                rank_idx: int,
                num_ranks: int,
                num_experts: int,
                num_topk: int,
                block_m: int) -> None:
    del num_topk
    num_experts_per_rank = num_experts // num_ranks
    sf_block_m = ((block_m + 127) // 128) * 128
    num_sms = torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count
    entries_by_expert, counts = build_expected_entries(all_topk_idx, rank_idx, num_experts_per_rank)

    torch.testing.assert_close(cumulative_stats.cpu(), counts.to(torch.int32).cpu(), rtol=0, atol=0)

    expert_status = buffer.expert_recv_count_sum[:num_experts_per_rank].detach().cpu().tolist()
    arrival = buffer.l1_arrival_count.detach().cpu()
    metadata = buffer.token_src_metadata.detach().cpu()

    pool_block_offset = 0
    local_start = rank_idx * num_experts_per_rank
    for local_expert, expected_entries in enumerate(entries_by_expert):
        count = len(expected_entries)
        status = int(expert_status[local_expert])
        assert (status & 0xffffffff) == count, (local_expert, status, count)
        assert (status >> 32) == num_sms * num_ranks, (local_expert, status, num_sms, num_ranks)

        expected_set = set(expected_entries)
        actual_set = set()
        base_pool_token = pool_block_offset * block_m
        expected_global_expert = local_start + local_expert

        for token_in_expert in range(count):
            pool_token_idx = base_pool_token + token_in_expert
            src_rank, src_token_idx, src_topk_idx = [int(v) for v in metadata[pool_token_idx].tolist()]
            actual_entry = (src_rank, src_token_idx, src_topk_idx)
            assert actual_entry in expected_set, (local_expert, token_in_expert, actual_entry, expected_set)
            assert actual_entry not in actual_set, (local_expert, token_in_expert, actual_entry)
            actual_set.add(actual_entry)

            expert_idx = int(all_topk_idx[src_rank, src_token_idx, src_topk_idx].item())
            assert expert_idx == expected_global_expert, (pool_token_idx, expert_idx, expected_global_expert)

            torch.testing.assert_close(
                buffer.l1_acts[pool_token_idx].to(torch.float32).cpu(),
                all_x_f32[src_rank, src_token_idx].cpu(),
                rtol=0, atol=0)
            torch.testing.assert_close(
                buffer.l1_topk_weights[pool_token_idx].cpu(),
                all_topk_weights[src_rank, src_token_idx, src_topk_idx].cpu(),
                rtol=0, atol=0)

            sf_pool_token_idx = (pool_block_offset + token_in_expert // block_m) * sf_block_m + token_in_expert % block_m
            torch.testing.assert_close(
                buffer.l1_acts_sf[sf_pool_token_idx].cpu(),
                all_x_sf[src_rank, src_token_idx].cpu(),
                rtol=0, atol=0)

        assert actual_set == expected_set, (local_expert, actual_set, expected_set)

        for block_idx in range(ceil_div(count, block_m)):
            expected_arrivals = min(block_m, count - block_idx * block_m)
            actual_arrivals = int(arrival[pool_block_offset + block_idx].item())
            assert actual_arrivals == expected_arrivals, (local_expert, block_idx, actual_arrivals, expected_arrivals)

        pool_block_offset += ceil_div(count, block_m)


def run_case(args: argparse.Namespace,
             group: dist.ProcessGroup,
             rank_idx: int,
             num_ranks: int,
             buffer: deep_gemm.SymmBuffer,
             weights,
             num_tokens: int) -> None:
    hidden = args.hidden
    num_topk = args.num_topk
    num_experts = args.num_experts

    x_f32 = torch.randn((num_tokens, hidden), dtype=torch.float32, device='cuda')
    x_fp8 = x_f32.to(torch.float8_e4m3fn)
    x_sf = torch.rand((num_tokens, hidden // 128), dtype=torch.float32, device='cuda') + 0.5
    topk_idx, topk_weights = make_topk(num_tokens, num_experts, num_topk, rank_idx)

    all_x_f32 = gather_same_shape(x_fp8.to(torch.float32), group)
    all_x_sf = gather_same_shape(x_sf, group)
    all_topk_idx = gather_same_shape(topk_idx, group)
    all_topk_weights = gather_same_shape(topk_weights, group)

    buffer.x[:num_tokens].copy_(x_fp8)
    buffer.x_sf[:num_tokens].copy_(x_sf)
    buffer.topk_idx[:num_tokens].copy_(topk_idx)
    buffer.topk_weights[:num_tokens].copy_(topk_weights)
    torch.cuda.synchronize()
    dist.barrier(group=group)

    cumulative_stats = torch.zeros((num_experts // num_ranks,), dtype=torch.int32, device='cuda')
    y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda')
    deep_gemm.fp8_mega_moe(y, weights[0], weights[1], buffer,
                           cumulative_local_expert_recv_stats=cumulative_stats)
    torch.cuda.synchronize()

    block_m = get_block_m(num_tokens, num_ranks, num_topk, num_experts)
    verify_case(buffer, all_x_f32, all_x_sf, all_topk_idx, all_topk_weights,
                cumulative_stats, rank_idx, num_ranks, num_experts, num_topk, block_m)
    dist.barrier(group=group)
    if rank_idx == 0:
        print(f'[PASSED] tokens={num_tokens}, block_m={block_m}', flush=True)


def main() -> None:
    parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 2 dispatch-only correctness')
    parser.add_argument('--tokens-list', type=str, default='0,8,48,192')
    parser.add_argument('--num-max-tokens-per-rank', type=int, default=384)
    parser.add_argument('--hidden', type=int, default=512)
    parser.add_argument('--intermediate-hidden', type=int, default=256)
    parser.add_argument('--num-experts', type=int, default=16)
    parser.add_argument('--num-topk', type=int, default=6)
    parser.add_argument('--local-rank', type=int, default=None)
    args = parser.parse_args()

    local_rank = args.local_rank if args.local_rank is not None else int(os.environ.get('LOCAL_RANK', '0'))
    rank_idx, num_ranks, group = init_test_dist(local_rank)
    assert torch.cuda.get_device_capability(torch.cuda.current_device())[0] == 9
    assert args.num_experts % num_ranks == 0
    assert args.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
    assert args.num_topk <= args.num_experts

    torch.manual_seed(1234 + rank_idx)
    random.seed(1234 + rank_idx)

    tokens_list = parse_tokens_list(args.tokens_list)
    max_tokens = max([args.num_max_tokens_per_rank] + tokens_list)
    buffer = deep_gemm.get_symm_buffer_for_mega_moe(
        group, args.num_experts, max_tokens, args.num_topk,
        args.hidden, args.intermediate_hidden)
    weights = make_weights(args.num_experts // num_ranks, args.hidden, args.intermediate_hidden)

    if rank_idx == 0:
        print(f'[Phase 2] ranks={num_ranks}, tokens_list={tokens_list}, '
              f'hidden={args.hidden}, intermediate={args.intermediate_hidden}, '
              f'experts={args.num_experts}, topk={args.num_topk}', flush=True)

    for num_tokens in tokens_list:
        run_case(args, group, rank_idx, num_ranks, buffer, weights, num_tokens)

    dist.barrier(group=group)
    buffer.destroy()
    dist.destroy_process_group()
    if rank_idx == 0:
        print('[PASSED] SM90 MegaMoE Phase 2 dispatch-only correctness', flush=True)


if __name__ == '__main__':
    main()
