// Package tokenizer 按模型估算 token 数（tiktoken-go）。
// 用于流式中断时对已生成内容做近似计费；未知模型回退 cl100k_base。
package tokenizer

import (
	"sync"

	"github.com/tiktoken-go/tokenizer"
)

// 编码器缓存：tokenizer.Get/ForModel 每次调用都新建 codec，其中包含 regexp2
// 正则编译（回溯引擎，编译昂贵）与 map 分配。实测短文本 6.9µs、18KB 文本
// 1.05ms，绝大部分是这套固定开销而非真正的分词。
//
// codec 构建后只读（vocabulary/splitRegexp 不再变），regexp2.Regexp 自身带锁，
// 因此缓存实例可并发复用。键是模型名：调用方模型名有限，缓存自然有界。
var codecs sync.Map // string → tokenizer.Codec

// Count 估算文本 token 数；空文本返回 0。
func Count(text, model string) int {
	if text == "" {
		return 0
	}
	enc := codecFor(model)
	if enc == nil {
		return 0
	}
	toks, _, err := enc.Encode(text)
	if err != nil {
		return 0
	}
	return len(toks)
}

// codecFor 取（必要时构建并缓存）该模型对应的 codec；未知模型回退 cl100k_base。
func codecFor(model string) tokenizer.Codec {
	if v, ok := codecs.Load(model); ok {
		return v.(tokenizer.Codec)
	}
	enc, err := tokenizer.ForModel(tokenizer.Model(model))
	if err != nil || enc == nil {
		enc = fallback()
	}
	if enc == nil {
		return nil
	}
	codecs.Store(model, enc)
	return enc
}

// fallback 取 cl100k_base 编码器（单独缓存，避免每次未知模型都重建）。
func fallback() tokenizer.Codec {
	const key = "\x00cl100k"
	if v, ok := codecs.Load(key); ok {
		return v.(tokenizer.Codec)
	}
	enc, err := tokenizer.Get(tokenizer.Cl100kBase)
	if err != nil || enc == nil {
		return nil
	}
	codecs.Store(key, enc)
	return enc
}