#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/cluster_primitive.h"

namespace xpu3 {
namespace plugin {

__global__ void ComputeOrderKernel(const int* seq_lens_this_time,
                                   const int* seq_lens_encoder,
                                   const int* base_model_seq_lens_this_time,
                                   const int* base_model_seq_lens_encoder,
                                   const int* accept_nums,
                                   int* position_map,
                                   int* output_token_num,
                                   const int bsz,
                                   const int actual_draft_token_num,
                                   const int input_token_num) {
  int tid = core_id() * cluster_num() + cluster_id();
  if (tid != 0) {
    return;
  }

  // 256 * int
  char lm[6 * 1024];
  int buf_size = 6 * 1024 / (6 * sizeof(int));
  int* lm_base_model_seq_lens_this_time = (int*)lm;
  int* lm_base_model_seq_lens_encoder =
      lm_base_model_seq_lens_this_time + buf_size;
  int* lm_seq_lens_this_time = lm_base_model_seq_lens_encoder + buf_size;
  int* lm_accept_nums = lm_seq_lens_this_time + buf_size;
  int* lm_seq_lens_encoder = lm_accept_nums + buf_size;
  int* lm_position_map = lm_seq_lens_encoder + buf_size;

  int in_offset = 0;
  int out_offset = 0;
  for (int i = 0; i < bsz; i += buf_size) {
    int64_t read_size = min(static_cast<int64_t>(bsz - i), buf_size);
    GM2LM(seq_lens_encoder + i, lm_seq_lens_encoder, read_size * sizeof(int));
    for (int j = 0; j < read_size; j++) {
      int cur_base_model_seq_lens_this_time =
          lm_base_model_seq_lens_this_time[j];
      int cur_base_model_seq_lens_encoder = lm_base_model_seq_lens_encoder[j];
      int cur_seq_lens_this_time = lm_seq_lens_this_time[j];
      int accept_num = lm_accept_nums[j];
      int cur_seq_lens_encoder = lm_seq_lens_encoder[j];
      // 1. eagle encoder. Base step=1
      if (cur_seq_lens_encoder > 0) {
        for (int k = 0; k < cur_seq_lens_encoder; k += buf_size) {
          int64_t write_size =
              min(static_cast<int64_t>(cur_seq_lens_encoder - k),
                  static_cast<int64_t>(buf_size));
          for (int l = 0; l < write_size; l++) {
            lm_position_map[l] = out_offset;
            out_offset++;
          }
          mfence_lm();
          LM2GM(lm_position_map,
                position_map + in_offset,
                write_size * sizeof(int));
          in_offset += write_size;
        }
        mfence_lm();
      }
    }
  }

  for (int i = 0; i < bsz; i += buf_size) {
    int64_t read_size = min(static_cast<int64_t>(bsz - i), buf_size);
    GM2LM_ASYNC(base_model_seq_lens_this_time + i,
                lm_base_model_seq_lens_this_time,
                read_size * sizeof(int));
    GM2LM_ASYNC(base_model_seq_lens_encoder + i,
                lm_base_model_seq_lens_encoder,
                read_size * sizeof(int));
    GM2LM_ASYNC(
        seq_lens_this_time + i, lm_seq_lens_this_time, read_size * sizeof(int));
    GM2LM_ASYNC(accept_nums + i, lm_accept_nums, read_size * sizeof(int));
    GM2LM(seq_lens_encoder + i, lm_seq_lens_encoder, read_size * sizeof(int));
    for (int j = 0; j < read_size; j++) {
      int cur_base_model_seq_lens_this_time =
          lm_base_model_seq_lens_this_time[j];
      int cur_base_model_seq_lens_encoder = lm_base_model_seq_lens_encoder[j];
      int cur_seq_lens_this_time = lm_seq_lens_this_time[j];
      int accept_num = lm_accept_nums[j];
      int cur_seq_lens_encoder = lm_seq_lens_encoder[j];
      // 1. eagle encoder. Base step=1
      if (cur_seq_lens_encoder > 0) {
        continue;
        // 2. Base model stop at last verify-step.
      } else if (cur_base_model_seq_lens_this_time != 0 &&
                 cur_seq_lens_this_time == 0) {
        in_offset += cur_base_model_seq_lens_this_time;
        // 4. stopped
      } else if (cur_base_model_seq_lens_this_time == 0 &&
                 cur_seq_lens_this_time == 0) {
        // nothing happens
      } else {
        // accept_num << buf_size, so do not need split
        for (int i = 0; i < accept_num; i++) {
          lm_position_map[i] = out_offset++;
        }
        mfence_lm();
        LM2GM(lm_position_map,
              position_map + in_offset,
              accept_num * sizeof(int));
        in_offset += cur_base_model_seq_lens_this_time;
        mfence_lm();
      }
    }
  }
  mfence_lm();
  LM2GM(&out_offset, output_token_num, sizeof(int));
}

}  // namespace plugin
}  // namespace xpu3
