#include "xpu/kernel/cluster.h"
#include "xpu/kernel/cluster_partition.h"
#include "xpu/kernel/cluster_primitive.h"
// #include <stdio.h>
// using namespace std;

#include "xpu/kernel/xtdk.h"
#include "xpu/kernel/xtdk_io.h"

namespace xpu3 {
namespace plugin {

template <typename T>
static inline __device__ bool is_in_end(const T id,
                                        const T* end_ids,
                                        const int length) {
  for (int i = 0; i < length; i++) {
    if (id == end_ids[i]) {
      return true;
    }
  }
  return false;
}

__global__ void limit_thinking_content_length_kernel_v1(
    int64_t* next_tokens,
    const int* max_think_lens,
    const int64_t* step_idx,
    const int64_t* eos_token_ids,
    int* limit_think_status,
    bool* stop_flags,
    const int64_t think_end_id,
    const int bs,
    const int eos_token_id_len) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  int nclusters = cluster_num();
  if (clusterid != 0) return;
  __simd__ __local__ int64_t eos_token_ids_lm[256];

  for (int i = cid; i < bs; i += ncores) {
    int max_think_len_lm;
    int limit_think_status_lm;
    int64_t next_token_lm;
    int64_t step_idx_lm;
    bool stop_flags_lm;
    GM2LM_ASYNC(next_tokens + i, &next_token_lm, sizeof(int64_t));
    GM2LM_ASYNC(step_idx + i, &step_idx_lm, sizeof(int64_t));
    GM2LM_ASYNC(max_think_lens + i, &max_think_len_lm, sizeof(int));
    GM2LM_ASYNC(stop_flags + i, &stop_flags_lm, sizeof(bool));
    GM2LM_ASYNC(
        eos_token_ids, eos_token_ids_lm, sizeof(int64_t) * eos_token_id_len);
    GM2LM(limit_think_status + i, &limit_think_status_lm, sizeof(int));

    // 如果该序列未启用思考功能，则直接返回，默认值为 -1，表示不限制思考长度
    if (max_think_len_lm < 0) continue;
    // 如果在回复阶段, 且已经触发停止标志, 则直接返回, 无需多余执行.
    if (limit_think_status_lm == 2 && stop_flags_lm) continue;

    // ======================= 思考阶段控制 =======================
    // 阶段 1: 仍在思考 (status == 0), 检查是否需要强制结束
    if (limit_think_status_lm < 1) {
      // 当开启思考长度控制时，检查是否超时
      if ((step_idx_lm >= max_think_len_lm) ||
          is_in_end(next_token_lm, eos_token_ids_lm, eos_token_id_len)) {
        // 强制将当前token替换为结束思考的token
        next_token_lm = think_end_id;
        // 将状态推进到 1, 表示 "正在结束思考"
        limit_think_status_lm = 1;
        if (stop_flags_lm) {
          stop_flags_lm = false;
          LM2GM(&stop_flags_lm, stop_flags + i, sizeof(bool));
        }
      }
    }

    // ======================= 思考结束处理 =======================
    // 阶段 2: 检查是否已满足结束思考的条件 (status < 2)
    // 这种情况会处理两种场景:
    // 1. status == 0: 模型自己生成了 think_end_id
    // 2. status == 1: 上一阶段强制注入了 think_end_id
    if (limit_think_status_lm < 2) {
      if (next_token_lm == think_end_id) {
        // 确认思考结束，将状态推进到 2 (响应阶段)
        limit_think_status_lm = 2;
      }
    }

    // 写回更新后的 token
    LM2GM_ASYNC(&next_token_lm, next_tokens + i, sizeof(int64_t));
    // 更新全局状态
    LM2GM(&limit_think_status_lm, limit_think_status + i, sizeof(int));
  }
}

}  // namespace plugin
}  // namespace xpu3
