#pragma once

#include <cutlass/cutlass.h>

#include <deep_gemm/common/math.cuh>
#include <deep_gemm/common/exception.cuh>
#include <deep_gemm/common/tma_utils.cuh>
#include <deep_gemm/common/types.cuh>
#include <deep_gemm/comm/barrier.cuh>
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include <deep_gemm/ptx/ld_st.cuh>
#include <deep_gemm/ptx/utils.cuh>
#include <deep_gemm/scheduler/sm90_mega_moe.cuh>

namespace deep_gemm {

template <
    uint32_t kNumMaxTokensPerRank,
    uint32_t kHidden, uint32_t kIntermediateHidden,
    uint32_t kNumExperts, uint32_t kNumTopk,
    uint32_t kNumExpertsPerWave,
    uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
    uint32_t kStoreBlockM,
    uint32_t kNumMaxPoolTokens,
    uint32_t kNumPaddedSFPoolTokens,
    uint32_t kNumStages,
    uint32_t kNumDispatchThreads, uint32_t kNumTMAThreads,
    uint32_t kNumMathThreads,
    bool kCooperativeMode, bool kUseNMajorL2,
    uint32_t kNumSMs, uint32_t kNumRanks,
    float kActivationClamp, bool kFastMath,
    uint32_t L1_SHAPE_N = kIntermediateHidden * 2,
    uint32_t L1_SHAPE_K = kHidden,
    uint32_t L2_SHAPE_N = kHidden,
    uint32_t L2_SHAPE_K = kIntermediateHidden,
    uint32_t kNumThreads = kNumDispatchThreads + kNumTMAThreads + kNumMathThreads,
    uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks>
CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void
sm90_fp8_mega_moe_impl(void* y,
                       int* cumulative_local_expert_recv_stats,
                       const uint32_t num_tokens,
                       const __grid_constant__ layout::SymBuffer<kNumRanks> sym_buffer,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights,
                       const void* l1_weights_sf,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf,
                       const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights,
                       const void* l2_weights_sf) {
    DG_STATIC_ASSERT(kNumThreads == 384, "SM90 MegaMoE expects 384 threads");
    DG_STATIC_ASSERT(BLOCK_N == 128, "SM90 MegaMoE expects BLOCK_N=128");
    DG_STATIC_ASSERT(BLOCK_K == 128, "SM90 MegaMoE expects BLOCK_K=128");
    DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks");

#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
    DG_STATIC_ASSERT(kNumDispatchThreads == 64, "SM90 dispatch-only path expects 64 dispatch threads");
    DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of top-k experts");
    DG_STATIC_ASSERT(kHidden % 128 == 0, "SM90 activation SF expects per-128K groups");
    DG_STATIC_ASSERT(kHidden % sizeof(uint4) == 0, "Token copy expects 16-byte alignment");

    constexpr uint32_t kNumDispatchWarps = kNumDispatchThreads / 32;
    constexpr uint32_t kNumTokensPerWarp = 32 / kNumTopk;
    constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk;
    constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, 128u);
    constexpr uint32_t kNumMaxPoolBlocks = kNumMaxPoolTokens / layout::kMinCandidateBlockM;
    constexpr uint32_t kNumTokenUint4 = kHidden / sizeof(uint4);
    constexpr uint32_t kNumSFValues = kHidden / 128;
    DG_STATIC_ASSERT(kNumTokensPerWarp > 0, "Invalid number of top-k experts");
    DG_STATIC_ASSERT(kNumPaddedSFPoolTokens % SF_BLOCK_M == 0, "Invalid padded SF pool size");

    // Only the first two warps participate in Phase 2 dispatch. TMA/math warps stay idle.
    const uint32_t thread_idx = threadIdx.x;
    if (thread_idx >= kNumDispatchThreads)
        return;

    const uint32_t sm_idx = blockIdx.x;
    const uint32_t warp_idx = thread_idx / 32;
    const uint32_t lane_idx = ptx::get_lane_idx();

    constexpr uint32_t kDispatchBarrierIdx = 0;
    constexpr uint32_t kDispatchGridSyncIndex = 0;
    constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 1;
    constexpr uint32_t kBeforeDispatchPullBarrierTag = 2;
    const auto dispatch_sync = []() {
        ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);
    };

    const auto workspace = layout::Workspace(
        sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk);

    constexpr auto fp8_token_layout = layout::Data(kHidden);
    constexpr auto fp8_sf_layout = layout::Data(kHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden);
    constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 128 * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(int64_t)), false);
    constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * static_cast<uint32_t>(sizeof(float)), false);
    constexpr auto l1_topk_weights_layout = layout::Data(static_cast<uint32_t>(sizeof(float)), false);

    const auto input_token_buffer = layout::Buffer(
        fp8_token_layout, 1, kNumMaxTokensPerRank,
        workspace.get_end_ptr());
    const auto input_sf_buffer = layout::Buffer(
        fp8_sf_layout, 1, kNumMaxTokensPerRank,
        input_token_buffer.get_end_ptr());
    const auto input_topk_idx_buffer = layout::Buffer(
        input_topk_idx_layout, 1, kNumMaxTokensPerRank,
        input_sf_buffer.get_end_ptr());
    const auto input_topk_weights_buffer = layout::Buffer(
        input_topk_weights_layout, 1, kNumMaxTokensPerRank,
        input_topk_idx_buffer.get_end_ptr());

    const auto l1_token_buffer = layout::Buffer(
        fp8_token_layout, 1, kNumMaxPoolTokens,
        input_topk_weights_buffer.get_end_ptr());
    const auto l1_sf_buffer = layout::Buffer(
        fp8_sf_layout, 1, kNumPaddedSFPoolTokens,
        l1_token_buffer.get_end_ptr());
    const auto l1_topk_weights_buffer = layout::Buffer(
        l1_topk_weights_layout, 1, kNumMaxPoolTokens,
        l1_sf_buffer.get_end_ptr());
    const auto l2_token_buffer = layout::Buffer(
        fp8_intermediate_token_layout, 1, kNumMaxPoolTokens,
        l1_topk_weights_buffer.get_end_ptr());
    const auto l2_sf_buffer = layout::Buffer(
        fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens,
        l2_token_buffer.get_end_ptr());
    (void)l2_sf_buffer;

    constexpr uint32_t kSharedMemoryAlignment = 1024;
    extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[];
    const auto smem_expert_count = reinterpret_cast<uint32_t*>(smem_buffer);

    // Clean local dispatch workspace from any previous call. Barrier state is intentionally
    // persistent and must not be reset because grid/NVLink barriers use phase counters.
    for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
        *workspace.get_expert_send_count_ptr(i) = 0;
    for (uint32_t i = thread_idx; i < kNumRanks * kNumExpertsPerRank; i += kNumDispatchThreads)
        *workspace.get_expert_recv_count_ptr(i / kNumExpertsPerRank, i % kNumExpertsPerRank) = 0;
    for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads)
        *workspace.get_expert_recv_count_sum_ptr(i) = 0;
    for (uint32_t i = thread_idx; i < kNumMaxPoolBlocks; i += kNumDispatchThreads) {
        *workspace.get_l1_arrival_count_ptr(i) = 0;
        *workspace.get_l2_arrival_mask_ptr(i) = 0;
    }
    comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
                         kDispatchGridSyncIndex, kAfterWorkspaceCleanBarrierTag>(
        workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync);

    for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads)
        smem_expert_count[i] = 0;
    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

    const auto read_topk_idx = [&](const auto& process) {
        #pragma unroll 1
        for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp;
             i < num_tokens;
             i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) {
            const uint32_t token_offset = lane_idx / kNumTopk;
            const uint32_t token_idx = i + token_offset;
            int expert_idx = -1;
            if (token_idx < num_tokens and lane_idx < kNumActivateLanes) {
                expert_idx = static_cast<int>(
                    __ldg(input_topk_idx_buffer.get_base_ptr<int64_t>() + i * kNumTopk + lane_idx));
                if (expert_idx >= 0 and expert_idx < static_cast<int>(kNumExperts))
                    process(i * kNumTopk + lane_idx, static_cast<uint32_t>(expert_idx));
            }
            __syncwarp();
        }
    };

    // Count local outgoing token-topk entries per global expert.
    read_topk_idx([&](const uint32_t&, const uint32_t& expert_idx) {
        atomicAdd_block(smem_expert_count + expert_idx, 1);
    });
    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

    // Convert per-SM counts into global per-rank offsets. High 32 bits count arrived SMs.
    for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
        const uint64_t send_value = (1ull << 32) | static_cast<uint64_t>(smem_expert_count[i]);
        smem_expert_count[i] = static_cast<uint32_t>(
            ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value));
    }
    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

    // Write source token-topk indices into destination ranks' local-expert tables.
    read_topk_idx([&](const uint32_t& token_topk_idx, const uint32_t& expert_idx) {
        const uint32_t dst_rank_idx = expert_idx / kNumExpertsPerRank;
        const uint32_t dst_local_expert_idx = expert_idx - dst_rank_idx * kNumExpertsPerRank;
        const uint32_t dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1);
        const auto dst_ptr = workspace.get_src_token_topk_idx_ptr(
            dst_local_expert_idx, sym_buffer.rank_idx, dst_slot_idx);
        *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx;
    });

    // Wait until all local SMs have finished filling local send-count/source-index data.
    comm::grid_sync<kNumSMs, kDispatchGridSyncIndex>(
        workspace, sm_idx, thread_idx, dispatch_sync);

    // Publish per-rank expert counts and summed ready/count words to destination ranks.
    if (sm_idx == 0) {
        for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) {
            const uint32_t dst_rank_idx = i / kNumExpertsPerRank;
            const uint32_t dst_local_expert_idx = i - dst_rank_idx * kNumExpertsPerRank;
            const uint64_t expert_status = *workspace.get_expert_send_count_ptr(i);
            *sym_buffer.map(
                workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx),
                dst_rank_idx) = expert_status & 0xffffffffu;
            ptx::atomic_add_sys(
                sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx),
                expert_status);
        }
    }
    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

    // All ranks must observe the complete count/source-index tables before pulling.
    comm::nvlink_barrier<kNumRanks, kNumSMs, kNumDispatchThreads,
                         kDispatchGridSyncIndex, kBeforeDispatchPullBarrierTag>(
        workspace, sym_buffer, sm_idx, thread_idx, dispatch_sync,
        /* sync_prologue */ false, /* sync_epilogue */ true);

    auto scheduler = sched::SM90MegaMoEScheduler<
        BLOCK_M, BLOCK_N, BLOCK_K,
        L1_SHAPE_N, L1_SHAPE_K,
        L2_SHAPE_N, L2_SHAPE_K,
        kNumExpertsPerRank,
        kNumExpertsPerWave,
        kNumSMs, kNumRanks,
        kUseNMajorL2>(workspace);
    scheduler.fetch_expert_recv_count();

    if (sm_idx == 0 and cumulative_local_expert_recv_stats != nullptr) {
        for (uint32_t i = thread_idx; i < kNumExpertsPerRank; i += kNumDispatchThreads) {
            const auto num_recv_tokens = static_cast<uint32_t>(*workspace.get_expert_recv_count_sum_ptr(i));
            ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast<int>(num_recv_tokens));
        }
    }
    ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx);

    constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u);
    int current_expert_idx = -1;
    uint32_t stored_rank_count[kNumRanksPerLane] = {};
    uint32_t expert_start_idx = 0, expert_end_idx = 0;
    uint32_t expert_pool_block_offset = 0;

    constexpr uint32_t kNumGlobalDispatchWarps = kNumSMs * kNumDispatchWarps;
    for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalDispatchWarps) {
        int old_expert_idx = current_expert_idx;
        while (token_idx >= expert_end_idx) {
            if (++ current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
                break;

            expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M);
            expert_start_idx = expert_end_idx;
            expert_end_idx += scheduler.get_num_tokens(static_cast<uint32_t>(current_expert_idx));
        }

        if (current_expert_idx >= static_cast<int>(kNumExpertsPerRank))
            break;

        if (old_expert_idx != current_expert_idx) {
            old_expert_idx = current_expert_idx;
            #pragma unroll
            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
                const uint32_t j = i * 32 + lane_idx;
                stored_rank_count[i] = j < kNumRanks ?
                    static_cast<uint32_t>(*workspace.get_expert_recv_count_ptr(j, static_cast<uint32_t>(current_expert_idx))) : 0;
            }
        }

        uint32_t current_rank_in_expert_idx = 0;
        uint32_t remaining[kNumRanksPerLane];
        #pragma unroll
        for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
            remaining[i] = stored_rank_count[i];

        uint32_t offset = 0;
        const uint32_t token_idx_in_expert = token_idx - expert_start_idx;
        uint32_t slot_idx = token_idx_in_expert;
        uint32_t token_idx_in_rank = 0;
        while (true) {
            uint32_t num_actives_in_lane = 0;
            uint32_t min_in_lane = 0xffffffffu;
            #pragma unroll
            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
                num_actives_in_lane += remaining[i] > 0;
                if (remaining[i] > 0)
                    min_in_lane = cute::min(min_in_lane, remaining[i]);
            }
            const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane);
            const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane);
            const uint32_t num_round_tokens = length * num_active_ranks;

            if (slot_idx < num_round_tokens) {
                const uint32_t slot_idx_in_round = slot_idx % num_active_ranks;
                uint32_t num_seen_ranks = 0;
                #pragma unroll
                for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) {
                    const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0);
                    const uint32_t num_active_lanes = __popc(mask);
                    if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes)
                        current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1);
                    num_seen_ranks += num_active_lanes;
                }
                token_idx_in_rank = offset + slot_idx / num_active_ranks;
                break;
            }

            slot_idx -= num_round_tokens;
            offset += length;
            #pragma unroll
            for (uint32_t i = 0; i < kNumRanksPerLane; ++ i)
                remaining[i] -= cute::min(remaining[i], length);
        }

        const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr(
            static_cast<uint32_t>(current_expert_idx), current_rank_in_expert_idx, token_idx_in_rank);
        const uint32_t src_token_idx = src_token_topk_idx / kNumTopk;
        const uint32_t src_topk_idx = src_token_topk_idx - src_token_idx * kNumTopk;
        const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert;
        const uint32_t sf_pool_token_idx =
            (expert_pool_block_offset + token_idx_in_expert / BLOCK_M) * SF_BLOCK_M +
            (token_idx_in_expert % BLOCK_M);

        const auto remote_token_ptr = sym_buffer.map(
            input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr<uint4>(),
            current_rank_in_expert_idx);
        const auto local_token_ptr = l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr<uint4>();
        #pragma unroll
        for (uint32_t i = lane_idx; i < kNumTokenUint4; i += 32)
            local_token_ptr[i] = remote_token_ptr[i];

        const auto remote_sf_ptr = sym_buffer.map(
            input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr<float>(),
            current_rank_in_expert_idx);
        const auto local_sf_ptr = l1_sf_buffer.get_base_ptr<float>();
        #pragma unroll
        for (uint32_t i = lane_idx; i < kNumSFValues; i += 32)
            local_sf_ptr[i * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[i];

        __syncwarp();
        if (lane_idx == 0) {
            const float weight = *sym_buffer.map(
                input_topk_weights_buffer.get_base_ptr<float>() + src_token_topk_idx,
                current_rank_in_expert_idx);
            *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr<float>() = weight;
            *workspace.get_token_src_metadata_ptr(pool_token_idx) = {
                current_rank_in_expert_idx, src_token_idx, src_topk_idx};
        }
        __syncwarp();
        __threadfence();
        __syncwarp();
        if (lane_idx == 0) {
            ptx::red_add_rel(
                workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1);
        }
        __syncwarp();
    }
#endif
    return;
}

} // namespace deep_gemm
