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

namespace fd_xpu3 {

__attribute__((global)) void recover_spec_decode_task(
    bool *stop_flags,
    int *seq_lens_this_time,
    int *seq_lens_encoder,
    int *seq_lens_decoder,
    int *step_seq_lens_decoder,
    int *block_tables,
    bool *is_block_step,
    int64_t *draft_tokens,
    const int64_t *step_draft_tokens,
    const int *step_seq_lens_this_time,
    const int bsz,
    const int block_num_per_seq,
    const int block_size,
    const int draft_tokens_len,
    const int num_extra_tokens) {
  int tid = core_id() * cluster_num() + cluster_id();
  int nthreads = cluster_num() * core_num();
  int start = -1;
  int end = -1;
  int output_padding_offset_lm;
  partition(tid, nthreads, static_cast<int>(bsz), 1, &start, &end);
  bool is_block_step_lm;
  int step_seq_lens_decoder_lm;
  int block_table_lm;
  bool value_false = false;
  int value_zero = 0;
  int step_seq_lens_this_time_lm;
  // max draft token = 6
  int64_t draft_tokens_lm[6];
  for (int i = start; i < end; i++) {
    GM2LM(is_block_step + i, &is_block_step_lm, sizeof(bool));
    if (is_block_step_lm == true) {
      GM2LM(step_seq_lens_decoder + i, &step_seq_lens_decoder_lm, sizeof(int));
      int max_possible_block_idx =
          (step_seq_lens_decoder_lm + num_extra_tokens) / block_size;
      max_possible_block_idx = min(max_possible_block_idx, block_num_per_seq);
      GM2LM(block_tables + max_possible_block_idx + i * block_num_per_seq,
            &block_table_lm,
            sizeof(int));
      if (block_table_lm != -1) {
        // can be recovered for decoding
        GM2LM(step_seq_lens_this_time + i,
              &step_seq_lens_this_time_lm,
              sizeof(int));

        LM2GM_ASYNC(&value_false, is_block_step + i, sizeof(bool));
        LM2GM_ASYNC(&value_false, stop_flags + i, sizeof(bool));
        LM2GM_ASYNC(&value_zero, seq_lens_encoder + i, sizeof(int));

        LM2GM_ASYNC(
            &step_seq_lens_this_time_lm, seq_lens_this_time + i, sizeof(int));
        LM2GM_ASYNC(
            &step_seq_lens_decoder_lm, seq_lens_decoder + i, sizeof(int));

        GM2LM(step_draft_tokens + i * draft_tokens_len,
              draft_tokens_lm,
              step_seq_lens_this_time_lm * sizeof(int64_t));
        LM2GM(draft_tokens_lm,
              draft_tokens + i * draft_tokens_len,
              step_seq_lens_this_time_lm * sizeof(int64_t));
      }
    }
  }
}

}  // namespace fd_xpu3
