#include "xpu/kernel/cluster_debug.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/xtdk.h"
#include "xpu/kernel/xtdk_math.h"
#include "xpu/kernel/xtdk_simd.h"

namespace fd_xpu3 {

static inline __device__ int v_reduce(int32x16_t &v0, int32x16_t &v1) {
  int res;
  v1 = vvadd_int32x16(v0, v1);
  auto v = vsrlp_int32x16(256, v1);
  v1 = vvadd_int32x16(v, v1);
  v = vsrlp_int32x16(128, v1);
  v1 = vvadd_int32x16(v, v1);
  v = vsrlp_int32x16(64, v1);
  v1 = vvadd_int32x16(v, v1);
  v = vsrlp_int32x16(32, v1);
  v1 = vvadd_int32x16(v, v1);
  res = vextract_int32x16(v1, 1);
  return res;
}
static inline __device__ int ClusterReduce(
    const _shared_ptr_ int *stop_flag_now_int_sm, int len) {
  int sum = 0;
  if (core_id() == 0) {
    int32x16_t vec_x_0;
    int32x16_t vec_x_1;
    int32x16_t vec_y_0 = vzero<int>();
    int32x16_t vec_y_1 = vzero<int>();
    for (int i = 0; i < len; i += 32) {
      vload2_sm(stop_flag_now_int_sm + i, vec_x_0, vec_x_1);
      vec_y_0 = vvadd_int32x16(vec_y_0, vec_x_0);
      vec_y_1 = vvadd_int32x16(vec_y_1, vec_x_1);
    }
    sum = v_reduce(vec_y_0, vec_y_1);
  }
  return sum;
}
__device__ bool is_in_end(const int64_t id,
                          __global_ptr__ const int64_t *end_ids,
                          int length) {
  bool flag = false;
  for (int i = 0; i < length; i++) {
    if (id == end_ids[i]) {
      return true;
    }
  }
  return flag;
}
__device__ inline bool is_in(__global_ptr__ const int64_t *candidates,
                             const int64_t draft,
                             const int candidate_len) {
  for (int i = 0; i < candidate_len; i++) {
    if (draft == candidates[i]) {
      return true;
    }
  }
  return false;
}

static __device__ inline unsigned int xorwow(unsigned int &state) {
  state ^= state >> 7;
  state ^= state << 9;
  state ^= state >> 13;
  return state;
}

__device__ int64_t
topp_sampling_kernel(__global_ptr__ const int64_t *candidate_ids,
                     __global_ptr__ const float *candidate_scores,
                     __global_ptr__ const float *dev_curand_states,
                     const int candidate_len,
                     const float topp) {
  const int tid = core_id();
  float sum_scores = 0.0f;
  float rand_top_p = *dev_curand_states * topp;
  // printf("debug rand_top_p:%f\n",rand_top_p);
  for (int i = 0; i < candidate_len; i++) {
    sum_scores += candidate_scores[i];
    if (rand_top_p <= sum_scores) {
      return candidate_ids[i];
    }
  }
  return candidate_ids[0];
}

// GREEDY / TARGET_MATCH: exact single-token match
__device__ inline bool verify_one_match(int64_t target_token,
                                        int64_t draft_token) {
  return target_token == draft_token;
}

__device__ inline bool verify_one_topp(
    __global_ptr__ const int64_t *verify_tokens_row,
    int64_t draft_token,
    int actual_cand_len) {
  return is_in(verify_tokens_row, draft_token, actual_cand_len);
}

// ============================================================
// VerifyContext — per-batch mutable state + accept helpers.
// Eliminates repeated EOS/max_dec_len check and output write
// patterns across Phase 1 and Phase 2.
// ============================================================
struct VerifyContext {
  // Immutable per-batch (set once at kernel entry)
  int bid;
  int max_step_tokens;
  int end_length;
  __global_ptr__ const int64_t *end_tokens;
  __global_ptr__ const int64_t *max_dec_len;
  __global_ptr__ const int64_t *step_input_ids_now;
  __global_ptr__ int64_t *step_output_ids;

  // Mutable per-batch state
  int64_t cur_step_idx;
  int output_len_now;
  bool stopped;

  // Emit a token at position `pos` to output in Phase 1.
  // Performs: step_idx check, EOS detection, token replacement, output write.
  // Returns true if this sequence should stop (EOS or max_dec_len hit).
  __device__ bool emit_token(int pos, int64_t token) {
    cur_step_idx++;
    bool is_eos = is_in_end(token, end_tokens, end_length);
    bool max_len_hit = (cur_step_idx >= max_dec_len[bid]);
    if (max_len_hit && !is_eos) {
      token = end_tokens[0];
    }
    step_output_ids[bid * max_step_tokens + pos] = token;
    output_len_now++;
    if (is_eos || max_len_hit) {
      stopped = true;
      return true;
    }
    return false;
  }

  // Emit the final token at position `pos` in Phase 2.
  // Same EOS/max_dec_len logic. Increments output_len_now since
  // Phase 2 produces one additional token.
  __device__ void emit_final_token(int pos, int64_t token) {
    cur_step_idx++;
    bool is_eos = is_in_end(token, end_tokens, end_length);
    bool max_len_hit = (cur_step_idx >= max_dec_len[bid]);
    if (max_len_hit && !is_eos) {
      token = end_tokens[0];
    }
    step_output_ids[bid * max_step_tokens + pos] = token;
    output_len_now++;
  }

  // TOPP-only: verify_window bulk-accept fallback.
  //
  // When draft token is NOT in top-p set but IS the top-2 token,
  // check verify_window consecutive positions for top-1 match.
  // If all match, bulk-accept from position i through ii.
  //
  // Returns the new loop position (i) after handling.
  // Sets *rejected=true if fallback was not triggered (caller should break).
  __device__ int try_verify_window_fallback(
      int i,
      bool *rejected,
      __global_ptr__ const int64_t *verify_tokens_now,
      int seq_len_this_time,
      int max_candidate_len,
      int verify_window) {
    int ii = i;
    if (max_candidate_len >= 2 &&
        verify_tokens_now[ii * max_candidate_len + 1] ==
            step_input_ids_now[ii + 1]) {
      // top-2 matches — scan verify_window consecutive top-1 matches
      int j = 0;
      ii += 1;
      for (; j < verify_window && ii < seq_len_this_time - 1; j++, ii++) {
        if (verify_tokens_now[ii * max_candidate_len] !=
            step_input_ids_now[ii + 1]) {
          break;
        }
      }
      if (j >= verify_window) {
        // Bulk accept all tokens from i to ii
        for (; i < ii; i++) {
          if (emit_token(i, step_input_ids_now[i + 1])) return i;
        }
        return i;  // continue outer loop from position ii
      }
    }
    // Fallback not triggered or insufficient window — reject
    *rejected = true;
    return i;
  }
};

__global__ void verify_draft_tokens(
    // Core I/O
    int64_t *step_output_ids,
    int *step_output_len,
    const int64_t *step_input_ids,  // draft tokens
    // Target model outputs (strategy-dependent interpretation)
    const int64_t
        *target_tokens,  // GREEDY:argmax, TARGET_MATCH:sampled, TOPP:unused
    // Candidate set for TOPP/GREEDY (TARGET_MATCH: unused)
    const int64_t *candidate_ids,
    const float *candidate_scores,
    const int *candidate_lens,
    // Sampling params
    const float *curand_states,  // nullptr for GREEDY/TARGET_MATCH
    const float *topp,
    // Metadata
    const bool *stop_flags,
    const int *seq_lens_encoder,
    const int *seq_lens_this_time,
    const int64_t *end_tokens,
    const bool *is_block_step,
    const int *cu_seqlens_q_output,
    const int *reasoning_status,
    // max_dec_len / step_idx for EOS/max-len detection (read-only)
    const int64_t *max_dec_len,
    const int64_t *step_idx,
    // Dimensions and config
    const int max_bsz,
    const int real_bsz,
    const int max_step_tokens,
    const int end_length,
    const int max_seq_len,
    const int max_candidate_len,
    const int verify_window,
    const int verify_strategy,  // 0=TOPP, 1=GREEDY, 2=TARGET_MATCH
    const bool reject_all,
    const bool accept_all) {
  const int64_t tid = core_id() * cluster_num() + cluster_id();
  const int64_t nthreads = cluster_num() * core_num();
  for (int64_t bid = tid; bid < real_bsz; bid += nthreads) {
    step_output_len[bid] = 0;
    if (bid >= real_bsz || is_block_step[bid] || stop_flags[bid]) continue;
    const int start_token_id = cu_seqlens_q_output[bid];
    // Pointers are strategy-dependent (may be nullptr for unused params)
    auto *candidate_ids_now =
        candidate_ids ? candidate_ids + start_token_id * max_candidate_len
                      : nullptr;
    auto *candidate_scores_now =
        candidate_scores ? candidate_scores + start_token_id * max_candidate_len
                         : nullptr;
    auto *candidate_lens_now =
        candidate_lens ? candidate_lens + start_token_id : nullptr;
    auto *target_tokens_now =
        target_tokens ? target_tokens + start_token_id : nullptr;

    // Initialize per-batch verification context
    VerifyContext ctx;
    ctx.bid = bid;
    ctx.max_step_tokens = max_step_tokens;
    ctx.end_length = end_length;
    ctx.end_tokens = end_tokens;
    ctx.max_dec_len = max_dec_len;
    ctx.step_input_ids_now = step_input_ids + bid * max_step_tokens;
    ctx.step_output_ids = step_output_ids;
    ctx.cur_step_idx = step_idx[bid];
    ctx.output_len_now = 0;
    ctx.stopped = false;

    // ======== Phase 1: Verify draft tokens ========
    int i = 0;
    for (; i < seq_lens_this_time[bid] - 1; i++) {
      // Early exit conditions: reject-all, prefill, reasoning
      if (reject_all || seq_lens_encoder[bid] != 0 ||
          reasoning_status[bid] == 1) {
        break;
      }

      // Accept-all override (debug/warmup)
      if (accept_all) {
        if (ctx.emit_token(i, ctx.step_input_ids_now[i + 1])) break;
        continue;
      }

      // Strategy dispatch
      bool accepted = false;
      switch (verify_strategy) {
        case 0: {  // TOPP
          auto actual_cand_len = candidate_lens_now[i] > max_candidate_len
                                     ? max_candidate_len
                                     : candidate_lens_now[i];
          accepted = verify_one_topp(candidate_ids_now + i * max_candidate_len,
                                     ctx.step_input_ids_now[i + 1],
                                     actual_cand_len);
          if (!accepted) {
            bool rejected = false;
            i = ctx.try_verify_window_fallback(i,
                                               &rejected,
                                               candidate_ids_now,
                                               seq_lens_this_time[bid],
                                               max_candidate_len,
                                               verify_window);
            if (ctx.stopped || rejected) goto phase1_done;
            continue;  // bulk accept succeeded, continue from new i
          }
          break;
        }
        case 1:  // GREEDY
        case 2:  // TARGET_MATCH
          accepted = verify_one_match(target_tokens_now[i],
                                      ctx.step_input_ids_now[i + 1]);
          break;
      }

      if (accepted) {
        if (ctx.emit_token(i, ctx.step_input_ids_now[i + 1])) break;
      } else {
        break;  // reject
      }
    }
  phase1_done:

    // ======== Phase 2: Output token for rejected/last position ========
    if (!ctx.stopped) {
      int64_t output_token;
      switch (verify_strategy) {
        case 0: {  // TOPP — stochastic sampling from candidate set
          auto actual_cand_len = candidate_lens_now[i] > max_candidate_len
                                     ? max_candidate_len
                                     : candidate_lens_now[i];
          output_token =
              topp_sampling_kernel(candidate_ids_now + i * max_candidate_len,
                                   candidate_scores_now + i * max_candidate_len,
                                   curand_states + bid,
                                   actual_cand_len,
                                   topp[bid]);
          break;
        }
        case 1:  // GREEDY — deterministic argmax from target_tokens
        case 2:  // TARGET_MATCH — target model's sampled token
          output_token = target_tokens_now[i];
          break;
      }
      ctx.emit_final_token(i, output_token);
    }
    step_output_len[bid] = ctx.output_len_now;
  }
}

}  // namespace fd_xpu3
