internal/context/token_cache.go

package context

import (
	"hash/fnv"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	log "github.com/sirupsen/logrus"
	"github.com/tiktoken-go/tokenizer"
)

// TokenCountCacheConfig holds configuration for token counting cache.
type TokenCountCacheConfig struct {
	Enabled    bool `yaml:"enabled" json:"enabled"`
	MaxEntries int  `yaml:"max-entries" json:"max_entries"`
	TTLMinutes int  `yaml:"ttl-minutes" json:"ttl_minutes"`
}

// DefaultTokenCountCacheConfig returns default configuration.
func DefaultTokenCountCacheConfig() *TokenCountCacheConfig {
	return &TokenCountCacheConfig{
		Enabled:    true,
		MaxEntries: 10000,
		TTLMinutes: 60,
	}
}

// cachedCount holds a cached token count with metadata.
type cachedCount struct {
	count     int
	createdAt time.Time
}

// TokenCountCache caches token counts for content to avoid repeated computation.
type TokenCountCache struct {
	cache      sync.Map // hash -> *cachedCount
	maxEntries int
	ttl        time.Duration
	size       atomic.Int64

	// Tokenizer cache
	tokenizers sync.Map // model -> tokenizer.Codec

	// Stats
	hits   atomic.Int64
	misses atomic.Int64

	mu sync.RWMutex
}

// NewTokenCountCache creates a new token count cache.
func NewTokenCountCache(maxEntries int) *TokenCountCache {
	if maxEntries <= 0 {
		maxEntries = 10000
	}

	cache := &TokenCountCache{
		maxEntries: maxEntries,
		ttl:        60 * time.Minute,
	}

	log.WithField("max_entries", maxEntries).Debug("token count cache initialized")
	return cache
}

// Count returns the token count for content, using cache if available.
func (tc *TokenCountCache) Count(model, content string) int {
	if content == "" {
		return 0
	}

	// Generate cache key
	key := tc.cacheKey(model, content)

	// Check cache
	if val, ok := tc.cache.Load(key); ok {
		cached := val.(*cachedCount)
		if time.Since(cached.createdAt) < tc.ttl {
			tc.hits.Add(1)
			return cached.count
		}
		// Expired, remove it
		tc.cache.Delete(key)
		tc.size.Add(-1)
	}

	tc.misses.Add(1)

	// Calculate token count
	count := tc.countTokens(model, content)

	// Cache the result (with eviction if needed)
	if tc.size.Load() >= int64(tc.maxEntries) {
		tc.evictOldest()
	}

	tc.cache.Store(key, &cachedCount{
		count:     count,
		createdAt: time.Now(),
	})
	tc.size.Add(1)

	return count
}

// CountBatch counts tokens for multiple contents efficiently.
func (tc *TokenCountCache) CountBatch(model string, contents []string) []int {
	results := make([]int, len(contents))
	for i, content := range contents {
		results[i] = tc.Count(model, content)
	}
	return results
}

// countTokens performs the actual token counting.
func (tc *TokenCountCache) countTokens(model, content string) int {
	codec := tc.getTokenizer(model)
	if codec == nil {
		// Fallback: rough estimate (1 token ≈ 4 characters for English)
		return len(content) / 4
	}

	count, err := codec.Count(content)
	if err != nil {
		// Fallback on error
		return len(content) / 4
	}

	// Apply model-specific adjustment
	adjustment := tc.getAdjustmentFactor(model)
	return int(float64(count) * adjustment)
}

// getTokenizer returns a tokenizer for the model.
func (tc *TokenCountCache) getTokenizer(model string) tokenizer.Codec {
	// Check cache
	if val, ok := tc.tokenizers.Load(model); ok {
		return val.(tokenizer.Codec)
	}

	// Create tokenizer
	var codec tokenizer.Codec
	var err error

	normalized := strings.ToLower(model)

	switch {
	case strings.Contains(normalized, "claude"):
		codec, err = tokenizer.Get(tokenizer.Cl100kBase)
	case strings.Contains(normalized, "gpt-5"):
		codec, err = tokenizer.ForModel(tokenizer.GPT5)
	case strings.Contains(normalized, "gpt-4o"):
		codec, err = tokenizer.ForModel(tokenizer.GPT4o)
	case strings.Contains(normalized, "gpt-4"):
		codec, err = tokenizer.ForModel(tokenizer.GPT4)
	case strings.Contains(normalized, "o1"), strings.Contains(normalized, "o3"):
		codec, err = tokenizer.ForModel(tokenizer.O1)
	default:
		codec, err = tokenizer.Get(tokenizer.O200kBase)
	}

	if err != nil {
		log.WithError(err).WithField("model", model).Debug("failed to get tokenizer")
		return nil
	}

	// Cache it
	actual, _ := tc.tokenizers.LoadOrStore(model, codec)
	return actual.(tokenizer.Codec)
}

// getAdjustmentFactor returns a multiplier for models where tiktoken may be inaccurate.
func (tc *TokenCountCache) getAdjustmentFactor(model string) float64 {
	normalized := strings.ToLower(model)

	// Claude models: tiktoken tends to underestimate
	if strings.Contains(normalized, "claude") {
		return 1.1
	}

	// Kiro uses Claude under the hood
	if strings.Contains(normalized, "kiro") {
		return 1.1
	}

	return 1.0
}

// cacheKey generates a unique key for model+content.
func (tc *TokenCountCache) cacheKey(model, content string) uint64 {
	h := fnv.New64a()
	h.Write([]byte(model))
	h.Write([]byte{0}) // Separator
	h.Write([]byte(content))
	return h.Sum64()
}

// evictOldest removes the oldest entries when cache is full.
func (tc *TokenCountCache) evictOldest() {
	// Simple eviction: remove ~10% of entries
	toRemove := tc.maxEntries / 10
	if toRemove < 1 {
		toRemove = 1
	}

	removed := 0
	now := time.Now()

	tc.cache.Range(func(key, value interface{}) bool {
		if removed >= toRemove {
			return false
		}

		cached := value.(*cachedCount)
		// Remove expired or oldest
		if now.Sub(cached.createdAt) > tc.ttl/2 {
			tc.cache.Delete(key)
			tc.size.Add(-1)
			removed++
		}
		return true
	})

	// If still need to remove more, just delete any
	if removed < toRemove {
		tc.cache.Range(func(key, value interface{}) bool {
			if removed >= toRemove {
				return false
			}
			tc.cache.Delete(key)
			tc.size.Add(-1)
			removed++
			return true
		})
	}
}

// Stats returns cache statistics.
type TokenCacheStats struct {
	Size    int64
	Hits    int64
	Misses  int64
	HitRate float64
}

// Stats returns current cache statistics.
func (tc *TokenCountCache) Stats() *TokenCacheStats {
	hits := tc.hits.Load()
	misses := tc.misses.Load()
	total := hits + misses

	var hitRate float64
	if total > 0 {
		hitRate = float64(hits) / float64(total) * 100
	}

	return &TokenCacheStats{
		Size:    tc.size.Load(),
		Hits:    hits,
		Misses:  misses,
		HitRate: hitRate,
	}
}

// Clear clears the cache.
func (tc *TokenCountCache) Clear() {
	tc.cache.Range(func(key, value interface{}) bool {
		tc.cache.Delete(key)
		return true
	})
	tc.size.Store(0)
	tc.hits.Store(0)
	tc.misses.Store(0)
}

// Preload warms the cache with common content.
func (tc *TokenCountCache) Preload(model string, contents []string) {
	for _, content := range contents {
		tc.Count(model, content)
	}
}