import type { ChatItemType } from '@fastgpt/global/core/chat/type.d';
import { ChatRoleEnum } from '@fastgpt/global/core/chat/constants';
import { countMessagesTokens, countPromptTokens } from '@fastgpt/global/common/string/tiktoken';
import { adaptRole_Chat2Message } from '@fastgpt/global/core/chat/adapt';

/* slice chat context by tokens */
export function ChatContextFilter({
  messages = [],
  maxTokens
}: {
  messages: ChatItemType[];
  maxTokens: number;
}) {
  if (!Array.isArray(messages)) {
    return [];
  }
  const rawTextLen = messages.reduce((sum, item) => sum + item.value.length, 0);

  // If the text length is less than half of the maximum token, no calculation is required
  if (rawTextLen < maxTokens * 0.5) {
    return messages;
  }

  // filter startWith system prompt
  const chatStartIndex = messages.findIndex((item) => item.obj !== ChatRoleEnum.System);
  const systemPrompts: ChatItemType[] = messages.slice(0, chatStartIndex);
  const chatPrompts: ChatItemType[] = messages.slice(chatStartIndex);

  // reduce token of systemPrompt
  maxTokens -= countMessagesTokens({
    messages: systemPrompts
  });

  // 根据 tokens 截断内容
  const chats: ChatItemType[] = [];

  // 从后往前截取对话内容
  for (let i = chatPrompts.length - 1; i >= 0; i--) {
    const item = chatPrompts[i];
    chats.unshift(item);

    const tokens = countPromptTokens(item.value, adaptRole_Chat2Message(item.obj));
    maxTokens -= tokens;

    /* 整体 tokens 超出范围, system必须保留 */
    if (maxTokens <= 0) {
      chats.shift();
      break;
    }
  }

  return [...systemPrompts, ...chats];
}