/*
 * Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#pragma once
#include <cuda_bf16.h>
#include <cuda_fp16.h>

namespace tensorrt_llm::kernels::moe_alltoall {

// Configuration constants
static constexpr int kMaxTopK = 22;     // Maximum supported top-k experts per token
static constexpr int kMaxPayloads = 6;  // Maximum number of different payload types
static constexpr int kMaxRanks = 64;    // Maximum supported EP size
// uint64 words needed to hold the active-rank bitmask (kRankMaskWords * 64 must be >= kMaxRanks).
static constexpr int kRankMaskWords = (kMaxRanks + 63) / 64;
static_assert(kRankMaskWords * 64 >= kMaxRanks, "active_rank_mask too small for kMaxRanks");

// Describes a single payload type to be communicated
struct PayloadDescriptor {
  void const* src_data;    // Source data pointer [local_num_tokens, elements_per_token]
  int element_size;        // Size of each element in bytes
  int elements_per_token;  // Number of elements per token (e.g., hidden_size, top_k)
};

// Kernel pointers packed into a struct for device access
// Dispatch kernel pointers - const source data
struct DispatchKernelPointers {
  // Payload pointers
  void const* src_data_ptrs[kMaxPayloads];      // Array of source data pointers
  void* recv_buffers[kMaxRanks][kMaxPayloads];  // 2D array of receive buffer pointers
  int payload_bytes_per_token[kMaxPayloads];    // Bytes per token for each payload

  // Completion flags for synchronization
  uint32_t*
      completion_flags[kMaxRanks];  // If completion_flags[target_rank][source_rank] == *flag_val,
                                    // then source rank has signaled the target rank
  uint32_t* flag_val;  // The value of the flag for this round (stored on the local rank)

  // Local aux data pointers
  int* send_counters;             // [ep_size] How many tokens have been sent to each target rank
  int* recv_counters[kMaxRanks];  // How many tokens have been received from each source rank. Each
                                  // rank has [ep_size] counters
  int* local_token_counter;       // Atomic counter for completed tokens

  // Top-K compact routing info per local token (size: [local_num_tokens, top_k])
  int* topk_target_ranks;  // target rank per k, -1 for duplicates
  int* topk_send_indices;  // dst index per k, -1 for duplicates

  // Optional: statistics for EPLB
  int const* eplb_local_stats;          // [eplb_stats_num_experts]
  int* eplb_gathered_stats[kMaxRanks];  // [ep_size, eplb_stats_num_experts] per rank

  // Active-rank bitmask: bit i set => rank i is alive and participates in this collective.
  // Tokens routed to a masked rank are dropped (topk_*[k] = -1); flag writes/waits to/from
  // masked peers are skipped. The local rank's own bit must always be set (checked at launch).
  uint64_t active_rank_mask[kRankMaskWords];
};

// Combine kernel pointers - non-const output in src_data_ptrs[0], const recv buffers
struct CombineKernelPointers {
  // Payload pointers
  void* src_data_ptrs[kMaxPayloads];                  // src_data_ptrs[0] is output
  void* output_scales;                                // Optional output scales
  void const* recv_buffers[kMaxRanks][kMaxPayloads];  // 2D array of receive buffer pointers (const)

  // Completion flags for synchronization
  uint32_t*
      completion_flags[kMaxRanks];  // If completion_flags[target_rank][source_rank] == *flag_val,
                                    // then source rank has signaled the target rank
  uint32_t* flag_val;  // The value of the flag for this round (stored on the local rank)

  // Top-K compact routing info per local token (size: [local_num_tokens, top_k])
  int const* topk_target_ranks;  // target rank per k, -1 for duplicates
  int const* topk_send_indices;  // dst index per k, -1 for duplicates

  // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Combine skips flag
  // writes/waits to/from masked peers. Per-token accumulation does not recheck this mask
  // (dispatch already dropped masked targets via topk_send_indices == -1); the mask is a
  // per-collective snapshot, not something that changes between a dispatch and its combine.
  uint64_t active_rank_mask[kRankMaskWords];
};

// Dispatch phase parameters
struct MoeA2ADispatchParams {
  bool enable_pdl;  // True: launch with programmatic dependent launch

  // EP configuration
  int ep_size;      // Number of EP ranks
  int ep_rank;      // Current EP rank
  int num_experts;  // Total number of experts (supports non-divisible num_experts % ep_size)

  // Token configuration
  int local_num_tokens;     // Number of tokens on this rank
  int max_tokens_per_rank;  // Maximum tokens per rank for pre-allocation TODO: Rename to
                            // runtime_max_tokens_per_rank
  int top_k;                // Number of experts per token

  // Expert routing information
  int32_t const* token_selected_experts;  // [local_num_tokens, top_k]

  // Generic payloads
  int num_payloads;                          // Number of different payload types
  PayloadDescriptor payloads[kMaxPayloads];  // Array of payload descriptors

  // Local aux data
  uint32_t* flag_val;        // The value of the flag for this round (stored on the local rank)
  int* local_token_counter;  // Atomic counter for completed tokens on this rank
  int* send_counters;        // [ep_size] atomic counters - tracks tokens sent to each target rank
  int* topk_target_ranks;    // Top-K compact routing info per local token (size: [local_num_tokens,
                             // top_k]), target rank per k, -1 for duplicates
  int* topk_send_indices;    // Top-K compact routing info per local token (size: [local_num_tokens,
                             // top_k]), dst index per k, -1 for duplicates

  // Distributed aux data and recv buffers
  int* recv_counters[kMaxRanks];  // tracks tokens received from each source rank. Each rank has
                                  // [ep_size] counters
  uint32_t*
      completion_flags[kMaxRanks];  // If completion_flags[target_rank][source_rank] == *flag_val,
                                    // then source rank has signaled the target rank
  void* recv_buffers[kMaxRanks][kMaxPayloads];  // Per-rank receive buffers for each payload

  // Optional: statistics for EPLB
  bool enable_eplb;                     // Whether to enable EPLB
  int eplb_stats_num_experts;           // Number of experts for EPLB stats
  int const* eplb_local_stats;          // [eplb_stats_num_experts]
  int* eplb_gathered_stats[kMaxRanks];  // [ep_size, eplb_stats_num_experts] per rank

  // Whether to instantiate a kernel variant that checks active_rank_mask at all. Defaults to
  // false, which compiles out every is_rank_active() branch (dead-rank routing, peer counter
  // mirroring, EPLB-stat gather, completion-flag write/wait) for the common no-fault-tolerance
  // case. Must be true to use active_rank_mask for anything other than all-ones.
  bool enable_rank_mask{false};

  // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Only consulted by the
  // kernel when enable_rank_mask is true; defaults to all-ones for backwards-compatible
  // "no masking" behavior. NOTE: this initializer list must have exactly kRankMaskWords
  // elements; update it alongside any change to kMaxRanks.
  uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}};

  // CUDA stream
  cudaStream_t stream;
};

// Dispatch kernels
void moe_a2a_dispatch_launch(MoeA2ADispatchParams const& params);
// Prepare for dispatch: zero send_counters, local_token_counter and increment flag_val
void moe_a2a_prepare_dispatch_launch(MoeA2ADispatchParams const& params);

enum class MoeA2ACombineQuantMode : uint32_t {
  NONE,
  MXFP8,
  MXFP4,
  NVFP4,
};

// NOTE(siyuan): keep this align with include/flashinfer/fp4_layout.cuh
enum class MoeA2ACombineSwizzleSFMode : uint32_t {
  SWIZZLE_128x4,
  SWIZZLE_8x4,
  LINEAR,
};

// Combine phase parameters
struct MoeA2ACombineParams {
  bool enable_pdl;  // True: launch with programmatic dependent launch

  // EP configuration
  int ep_size;  // Number of EP ranks
  int ep_rank;  // Current EP rank

  // Token configuration
  int local_num_tokens;     // Number of tokens on this rank
  int max_tokens_per_rank;  // Maximum tokens per rank for pre-allocation TODO: Rename to
                            // runtime_max_tokens_per_rank
  int top_k;                // Number of experts per token

  // If true, recv buffers contain FP8 data (either pre-staged or quantized in-place by prepare).
  // The combine kernel reads FP8 and writes BF16 output.
  bool use_low_precision;

  // Prepare-only field: original payload tensor pointer used to stage into workspace
  void const* prepare_payload;

  // Output tensor
  MoeA2ACombineQuantMode quant_mode{MoeA2ACombineQuantMode::NONE};  // Output quantization mode
  MoeA2ACombineSwizzleSFMode swizzle_mode{
      MoeA2ACombineSwizzleSFMode::LINEAR};  // Output swizzle mode
  void* output_data;                        // Output buffer [local_num_tokens, elements_per_token]
  void* output_scales;                      // Optional output scales for quantized outputs
  float output_scalar_scale{1.0f};  // Per-tensor global scale applied before FP4 block scaling
                                    // (SFScaleVal); ignored by MXFP8/MXFP4 paths

  // Payload information
  int elements_per_token;    // Number of elements per token
  nvinfer1::DataType dtype;  // Data type for proper summation

  // Local aux data
  uint32_t* flag_val;        // The value of the flag for this round (stored on the local rank)
  int* topk_target_ranks;    // Top-K compact routing info per local token (size: [local_num_tokens,
                             // top_k]), target rank per k, -1 for duplicates
  int* topk_send_indices;    // Top-K compact routing info per local token (size: [local_num_tokens,
                             // top_k]), dst index per k, -1 for duplicates
  int const* recv_counters;  // [ep_size] number of valid tokens per source rank for this target

  // Distributed aux data and recv buffers
  uint32_t*
      completion_flags[kMaxRanks];  // If completion_flags[target_rank][source_rank] == *flag_val,
                                    // then source rank has signaled the target rank
  void const* recv_buffers[kMaxRanks];  // Per-rank receive buffers (only for single payload)

  // Whether to instantiate a kernel variant that checks active_rank_mask in peer
  // synchronization. Defaults to false, which compiles out the is_rank_active() branches
  // (completion-flag write/wait) for the common no-fault-tolerance case.
  bool enable_rank_mask{false};

  // Active-rank bitmask: see DispatchKernelPointers::active_rank_mask. Only consulted by the
  // kernel when enable_rank_mask is true; defaults to all-ones for backwards-compatible
  // "no masking" behavior. NOTE: this initializer list must have exactly kRankMaskWords
  // elements; update it alongside any change to kMaxRanks.
  uint64_t active_rank_mask[kRankMaskWords] = {~uint64_t{0}};

  // CUDA stream
  cudaStream_t stream;
};

// Combine kernels
void moe_a2a_combine_launch(MoeA2ACombineParams const& params);

void moe_a2a_prepare_combine_launch(MoeA2ACombineParams const& params);

// Sanitize expert IDs for invalid tokens
// expert_ids: [ep_size, max_tokens_per_rank, top_k] (int32)
// recv_counters: [ep_size] (int32), number of valid tokens per source
// invalid_id: value to fill for invalid tokens' expert ids
void moe_a2a_sanitize_expert_ids_launch(int32_t* expert_ids, int32_t const* recv_counters,
                                        int32_t invalid_id, int ep_size, int max_tokens_per_rank,
                                        int top_k, cudaStream_t stream, bool enable_pdl);

}  // namespace tensorrt_llm::kernels::moe_alltoall
