import argparse
import inspect
import os
import pathlib
import sys
from typing import 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


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 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, (intermediate_hidden * 2 + 127) // 128, hidden // 128),
        dtype=torch.float32, device='cuda')
    l2_weights_sf = torch.ones(
        (num_experts_per_rank, (hidden + 127) // 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 main() -> None:
    parser = argparse.ArgumentParser(description='SM90 MegaMoE Phase 1 interface smoke')
    parser.add_argument('--num-tokens', type=int, default=8)
    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.hidden % 128 == 0 and args.intermediate_hidden % 128 == 0
    assert args.num_experts % num_ranks == 0

    buffer = deep_gemm.get_symm_buffer_for_mega_moe(
        group, args.num_experts,
        args.num_max_tokens_per_rank, args.num_topk,
        args.hidden, args.intermediate_hidden)
    assert buffer.x.shape == (buffer.num_max_tokens_per_rank, args.hidden)
    assert buffer.x_sf.shape == (buffer.num_max_tokens_per_rank, args.hidden // 128)
    assert buffer.x_sf.dtype == torch.float32
    assert buffer.l1_acts.shape[1] == args.hidden
    assert buffer.l1_acts_sf.shape[1] == args.hidden // 128
    assert buffer.l1_acts_sf.dtype == torch.float32
    assert buffer.l2_acts.shape[1] == args.intermediate_hidden
    assert buffer.l2_acts_sf.shape[1] == args.intermediate_hidden // 128
    assert buffer.l2_acts_sf.dtype == torch.float32

    num_tokens = args.num_tokens
    buffer.x[:num_tokens].copy_(torch.randn((num_tokens, args.hidden), device='cuda').to(torch.float8_e4m3fn))
    buffer.x_sf[:num_tokens].fill_(1.0)
    buffer.topk_idx[:num_tokens].fill_(rank_idx * (args.num_experts // num_ranks))
    buffer.topk_weights[:num_tokens].fill_(1.0)

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

    dist.barrier(group=group)
    buffer.destroy()
    dist.destroy_process_group()
    if rank_idx == 0:
        print('[PASSED] SM90 MegaMoE Phase 1 interface smoke', flush=True)


if __name__ == '__main__':
    main()