package tokenizer

import (
	"fmt"
	"strings"

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

// Count 计算字符串的 token 数量
func Count(text, model string) int {
	tkm, err := tiktoken.EncodingForModel(model)
	if err != nil {
		tkm, _ = tiktoken.GetEncoding("cl100k_base")
	}
	return len(tkm.Encode(text, nil, nil))
}

// Cost 计算模型调用成本（USD，按每百万 token 定价）
func Cost(model string, inputTokens, outputTokens int) float64 {
	var inputPrice, outputPrice float64

	switch {
	case strings.Contains(model, "gpt-4o-mini"):
		inputPrice = 0.15
		outputPrice = 0.60
	case strings.Contains(model, "gpt-4o"):
		inputPrice = 2.50
		outputPrice = 10.00
	case strings.Contains(model, "gpt-4-turbo"):
		inputPrice = 10.00
		outputPrice = 30.00
	case strings.Contains(model, "gpt-4"):
		inputPrice = 30.00
		outputPrice = 60.00
	case strings.Contains(model, "gpt-3.5-turbo"):
		inputPrice = 0.50
		outputPrice = 1.50
	case strings.Contains(model, "claude-3-5-sonnet"):
		inputPrice = 3.00
		outputPrice = 15.00
	case strings.Contains(model, "claude-3-opus"):
		inputPrice = 15.00
		outputPrice = 75.00
	case strings.Contains(model, "claude-3-haiku"):
		inputPrice = 0.25
		outputPrice = 1.25
	case strings.Contains(model, "claude"):
		inputPrice = 8.00
		outputPrice = 24.00
	case strings.Contains(model, "gemini-1.5-pro"):
		inputPrice = 3.50
		outputPrice = 10.50
	case strings.Contains(model, "gemini-1.5-flash"):
		inputPrice = 0.35
		outputPrice = 0.53
	case strings.Contains(model, "gemini"):
		inputPrice = 0.50
		outputPrice = 1.50
	default:
		inputPrice = 0.15
		outputPrice = 0.60
	}

	cost := float64(inputTokens)/1e6*inputPrice + float64(outputTokens)/1e6*outputPrice
	if cost < 0.000001 {
		cost = 0.000001
	}
	return cost
}

// CostWithModel 从数据库模型记录获取定价
func CostWithModel(inputTokens, outputTokens, cacheReadTokens, cacheCreationTokens int64, inputPrice, outputPrice, cacheReadPrice float64) float64 {
	cost := float64(inputTokens)/1e6*inputPrice +
		float64(outputTokens)/1e6*outputPrice +
		float64(cacheReadTokens)/1e6*cacheReadPrice +
		float64(cacheCreationTokens)/1e6*inputPrice*1.25
	if cost < 0.000001 {
		cost = 0.000001
	}
	return cost
}

func init() {
	_ = fmt.Sprintf // ensure fmt is used
}