#pragma once

#include <algorithm>
#include <unordered_set>

#include <deep_gemm/layout/mega_moe.cuh>

#include "../../utils/exception.hpp"
#include "../../utils/math.hpp"
#include "../../utils/system.hpp"

namespace deep_gemm {

struct SM90MegaMoEConfig {
    // Block tiling
    int block_m, block_n, block_k;
    int store_block_m;

    // Pool capacity and SF-padded token count
    int num_max_pool_tokens;
    int num_padded_sf_pool_tokens;

    // Number of experts to process per wave
    int num_experts_per_wave;

    // Pipeline stages and shared memory
    int num_stages, smem_size;

    // Thread layout (384 total: 64 dispatch + 64 TMA + 256 math)
    int num_dispatch_threads, num_tma_threads, num_math_threads;

    // Mode flags
    bool cooperative;
    bool use_n_major_l2;

    friend std::ostream& operator << (std::ostream& os, const SM90MegaMoEConfig& config) {
        os << "SM90MegaMoEConfig("
           << "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k
           << ", store_block_m=" << config.store_block_m
           << ", num_max_pool_tokens=" << config.num_max_pool_tokens
           << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens
           << ", num_experts_per_wave=" << config.num_experts_per_wave
           << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size
           << ", num_dispatch_threads=" << config.num_dispatch_threads
           << ", num_tma_threads=" << config.num_tma_threads
           << ", num_math_threads=" << config.num_math_threads
           << ", cooperative=" << config.cooperative
           << ", use_n_major_l2=" << config.use_n_major_l2 << ")";
        return os;
    }
};

static std::tuple<int, int, int, bool> get_block_config_for_sm90_mega_moe(
    const int& num_ranks, const int& num_experts,
    const int& num_max_tokens_per_rank, const int& num_topk,
    const int& num_tokens) {

    float num_expected_tokens_per_expert =
        static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;

    if (num_expected_tokens_per_expert <= 16.5) {
        // Extreme decode: RL long-tail, large EP
        return {32, 16, 256, false};
    } else if (num_expected_tokens_per_expert <= 64.5) {
        // Medium decode
        return {64, 32, 256, false};
    } else {
        // Large M prefill / large EP decode
        return {128, 32, 256, true};
    }
}

static int get_num_experts_per_wave_for_sm90_mega_moe(
    const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
    const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) {

    float expected_tokens_per_expert = static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
    if (expected_tokens_per_expert < 1) {
        return num_experts_per_rank;
    }

    constexpr int kImbalanceFactor = 2;

    const int num_m_blocks = ceil_div(static_cast<int>(std::ceil(expected_tokens_per_expert)), block_m);
    const int num_n_blocks = (2 * intermediate_hidden) / block_n;
    const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks;

    int num_experts_per_wave = num_l1_blocks_per_expert > 0
        ? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1;
    num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank);

    while (num_experts_per_wave < num_experts_per_rank and num_experts_per_rank % num_experts_per_wave != 0)
        ++ num_experts_per_wave;

    return num_experts_per_wave;
}

static std::pair<int, int> get_pipeline_config_for_sm90_mega_moe(
    const int& smem_capacity,
    const int& num_experts, const int& hidden,
    const int& block_m, const int& block_n, const int& block_k, const int& store_block_m,
    const int& num_dispatch_threads, const int& num_math_threads,
    const bool& cooperative) {
    constexpr int kSmemAlignment = 1024;
    constexpr int kNumTMAStoreStages = 2;

    const int num_dispatch_warps = num_dispatch_threads / 32;
    const int num_math_warps = num_math_threads / 32;

    // Dispatch region
    const int smem_expert_count_size = align(
        num_experts * static_cast<int>(sizeof(uint32_t)), kSmemAlignment);
    const int smem_send_buffers_size = align(
        static_cast<int>(layout::Buffer(layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()),
        kSmemAlignment);
    const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size;

    // C/D output region: max of L1 FP8 (2 TMA stages) and L2 BF16 staging
    // L1: store_block_m * (block_n / 2) * kNumTMAStoreStages (FP8 = 1 byte)
    // L2: block_m * block_n * sizeof(BF16) (BF16 = 2 bytes)
    const int num_math_warpgroups = cooperative ? 2 : 1;
    const int smem_cd_l1 = num_math_warpgroups * store_block_m * (block_n / 2) * kNumTMAStoreStages;
    const int smem_cd_l2 = block_m * block_n * static_cast<int>(sizeof(nv_bfloat16));
    const int smem_cd = std::max(smem_cd_l1, smem_cd_l2);

    // Barriers: dispatch + full/empty pipeline (2 per stage) + combine (2 per math warp)
    const int smem_barriers = (num_dispatch_warps + 2 * 2 + num_math_warps * 2) * 8;

    // Amax reduction
    const int smem_amax_reduction = store_block_m * num_math_warps * static_cast<int>(sizeof(float));

    // Float SF per stage: align(2 * BLOCK_M * sizeof(float), 128)
    const int smem_sfa_per_stage = align(2 * block_m * static_cast<int>(sizeof(float)), 128);

    // Per-stage: A tile + B tile + SFA tile + full/empty barriers
    const int smem_per_stage = block_m * block_k + block_n * block_k + smem_sfa_per_stage + 2 * 8;

    // Fixed total
    const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers;

    const int num_stages = (smem_capacity - smem_fixed) / smem_per_stage;
    DG_HOST_ASSERT(num_stages >= 3);

    return {num_stages, smem_fixed + num_stages * smem_per_stage};
}

static SM90MegaMoEConfig get_sm90_mega_moe_config(
    const int& num_ranks, const int& num_experts, const int& num_experts_per_rank,
    const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk,
    const int& hidden, const int& intermediate_hidden,
    const int& num_padded_sf_pool_tokens) {

    const auto [block_m, store_block_m, num_math_threads, cooperative] =
        get_block_config_for_sm90_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens);
    const int block_n = 128;
    const int block_k = 128;

    const int num_max_pool_tokens = layout::get_num_max_pool_tokens(
        num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank);

    // Thread layout: 64 dispatch + 64 TMA + 256 math/epilogue = 384
    const int num_dispatch_threads = 64;
    const int num_tma_threads = 64;

    // Auto N-major L2: enabled when large M (high tokens per expert)
    const bool use_n_major_l2 = [&]() {
        auto env_val = get_env<int>("DG_SM90_MOE_NMAJOR", -1);
        if (env_val != -1)
            return env_val > 0;
        float expected = static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
        return expected >= 256;
    }();

    // Waves
    const int num_sms = device_runtime->get_num_sms();
    const int num_experts_per_wave = get_num_experts_per_wave_for_sm90_mega_moe(
        num_experts_per_rank, num_tokens, num_topk,
        intermediate_hidden, block_m, block_n, num_sms);

    // Pipeline
    constexpr int smem_capacity = 232448;
    const auto [num_stages, smem_size] = get_pipeline_config_for_sm90_mega_moe(
        smem_capacity,
        num_experts, hidden,
        block_m, block_n, block_k, store_block_m,
        num_dispatch_threads, num_math_threads,
        cooperative);

    const auto config = SM90MegaMoEConfig {
        block_m, block_n, block_k,
        store_block_m,
        num_max_pool_tokens, num_padded_sf_pool_tokens,
        num_experts_per_wave,
        num_stages, smem_size,
        num_dispatch_threads, num_tma_threads, num_math_threads,
        cooperative, use_n_major_l2
    };

    if (get_env<int>("DG_JIT_DEBUG") or get_env<int>("DG_PRINT_CONFIGS")) {
        const auto key = fmt::format(
            "SM90MegaMoEConfig(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})",
            num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk);
        static std::unordered_set<std::string> printed;
        if (printed.count(key) == 0) {
            std::cout << key << ": " << config << std::endl;
            printed.insert(key);
        }
    }
    return config;
}

} // namespace deep_gemm