import type { ModelMessage } from 'ai';

/**
 * Token 计数器
 * 使用简单的字符估算，不依赖外部库
 * 估算规则：
 * - 中文字符：约 1.5 字符/token
 * - 英文/数字：约 4 字符/token
 * - 混合内容取平均
 */
export class TokenCounter {
  /**
   * 估算文本的 token 数量
   */
  static estimateText(text: string): number {
    if (!text) return 0;

    // 统计中文字符数量
    const chineseChars = (text.match(/[\u4e00-\u9fff]/g) || []).length;
    // 其他字符数量
    const otherChars = text.length - chineseChars;

    // 中文约 1.5 字符/token，其他约 4 字符/token
    const chineseTokens = chineseChars / 1.5;
    const otherTokens = otherChars / 4;

    return Math.ceil(chineseTokens + otherTokens);
  }

  /**
   * 估算消息内容的 token 数量
   */
  static estimateContent(content: ModelMessage['content']): number {
    if (typeof content === 'string') {
      return this.estimateText(content);
    }

    if (Array.isArray(content)) {
      let total = 0;
      for (const part of content) {
        if (typeof part === 'string') {
          total += this.estimateText(part);
        } else if ('text' in part && typeof part.text === 'string') {
          total += this.estimateText(part.text);
        } else if ('result' in part) {
          // tool-result
          total += this.estimateText(JSON.stringify(part.result));
        } else if ('args' in part) {
          // tool-call
          total += this.estimateText(JSON.stringify(part.args));
          total += 20; // 工具名称等开销
        } else {
          // 其他类型，序列化估算
          total += this.estimateText(JSON.stringify(part));
        }
      }
      return total;
    }

    return 0;
  }

  /**
   * 估算单条消息的 token 数量
   */
  static estimateMessage(message: ModelMessage): number {
    let tokens = 0;

    // 角色标记开销
    tokens += 4;

    // 内容
    tokens += this.estimateContent(message.content);

    return tokens;
  }

  /**
   * 估算消息数组的总 token 数量
   */
  static estimateMessages(messages: ModelMessage[]): number {
    let total = 0;
    for (const message of messages) {
      total += this.estimateMessage(message);
    }
    // 消息间分隔开销
    total += messages.length * 3;
    return total;
  }

  /**
   * 格式化 token 数量显示
   */
  static format(tokens: number): string {
    if (tokens >= 1000) {
      return `${(tokens / 1000).toFixed(1)}k`;
    }
    return `${tokens}`;
  }
}
