#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 {

__global__ void limit_thinking_content_length_kernel_v2(
    int64_t* next_tokens,
    const int* max_think_lens,
    const int64_t* step_idx,
    int* limit_think_status,
    const bool* stop_flags,
    const int64_t think_end_id,
    const int64_t line_break_id,
    const int bs) {
  int cid = core_id();
  int ncores = core_num();
  int clusterid = cluster_id();
  int nclusters = cluster_num();
  if (clusterid != 0) return;

  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(stop_flags + i, &stop_flags_lm, sizeof(bool));
    GM2LM_ASYNC(max_think_lens + i, &max_think_len_lm, sizeof(int));
    GM2LM(limit_think_status + i, &limit_think_status_lm, sizeof(int));

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

    // ======================= 思考阶段控制 =======================
    // 阶段 1: 仍在思考 (status == 0), 检查是否需要强制结束
    // 阶段 2: 在替换 (status == 1), 检查是否替换结束
    if (limit_think_status_lm <= 1) {
      // 当开启思考长度控制时，检查是否超时
      if (step_idx_lm == max_think_len_lm) {
        // 强制将当前token替换为结束思考的token
        next_token_lm = line_break_id;
        limit_think_status_lm = 1;
      } else if (step_idx_lm == max_think_len_lm + 1) {
        // 强制将当前token替换为结束思考的token
        next_token_lm = think_end_id;
        limit_think_status_lm = 1;
      } else if (step_idx_lm == max_think_len_lm + 2) {
        // 强制将当前token替换为结束思考的token
        next_token_lm = line_break_id;
        limit_think_status_lm = 1;
      } else if (step_idx_lm == max_think_len_lm + 2) {
        // 强制将当前token替换为结束思考的token
        next_token_lm = line_break_id;
        limit_think_status_lm = 2;
      }
    }

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

    // 写回更新后的 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
