package context

import (
	"strings"
	"unicode"
)

// TokenEstimator estimates token counts for text using a simple heuristic.
// For production use, replace with a proper tokenizer (tiktoken, etc.).
type TokenEstimator struct {
	charsPerToken float64
}

// NewTokenEstimator creates a new estimator with the default ratio.
// English text averages ~4 chars/token, Chinese ~1.5 chars/token.
func NewTokenEstimator() *TokenEstimator {
	return &TokenEstimator{charsPerToken: 3.0}
}

// EstimateText estimates token count for a given text.
func (e *TokenEstimator) EstimateText(text string) int {
	if text == "" {
		return 0
	}

	// Count CJK characters as individual tokens
	cjkCount := 0
	nonCJKChars := 0
	for _, r := range text {
		if unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hiragana, r) || unicode.Is(unicode.Katakana, r) || unicode.Is(unicode.Hangul, r) {
			cjkCount++
		} else {
			nonCJKChars++
		}
	}

	// Non-CJK: estimate by chars/token ratio
	nonCJKTokens := int(float64(nonCJKChars) / e.charsPerToken)
	if nonCJKChars > 0 && nonCJKTokens == 0 {
		nonCJKTokens = 1
	}

	return cjkCount + nonCJKTokens
}

// EstimateMessage estimates token count for a single message (including role overhead).
func (e *TokenEstimator) EstimateMessage(msg interface{ GetRole() string; GetContent() string }) int {
	role := msg.GetRole()
	content := msg.GetContent()
	// Role tokens: ~1-2 tokens for role name
	roleTokens := len(strings.Fields(role)) + 1
	return roleTokens + e.EstimateText(content)
}

// EstimateMessages estimates total token count for a list of messages.
func (e *TokenEstimator) EstimateMessages(messages []Message) int {
	total := 0
	for _, m := range messages {
		total += e.EstimateText(m.Role) + e.EstimateText(m.Content) + 4 // role + content + formatting overhead
	}
	return total
}

// Message is a simplified message structure for estimation.
type Message struct {
	Role    string
	Content string
}

func (m Message) GetRole() string    { return m.Role }
func (m Message) GetContent() string { return m.Content }

// EstimateKVCache estimates the KV cache memory usage in bytes.
// Formula: input_tokens × layers × 2 (K+V) × hidden_dim × bytes_per_element
func EstimateKVCache(inputTokens, layers, hiddenDim, bytesPerElement int) int64 {
	return int64(inputTokens) * int64(layers) * 2 * int64(hiddenDim) * int64(bytesPerElement)
}

// EstimateKVCachePerToken estimates KV cache per token in bytes.
func EstimateKVCachePerToken(layers, hiddenDim, bytesPerElement int) int64 {
	return int64(layers) * 2 * int64(hiddenDim) * int64(bytesPerElement)
}