package tokeniser

import (
	"fmt"
	"strings"
	"sync"

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

// Encoding names for different model families
const (
	EncodingCl100kBase = "cl100k_base" // GPT-4, GPT-3.5-turbo, embeddings
	EncodingO200kBase  = "o200k_base"  // GPT-4o, GPT-4.1, GPT-4.5
	EncodingP50kBase   = "p50k_base"   // Codex models, davinci-002, davinci-003
	EncodingR50kBase   = "r50k_base"   // GPT-3 models
)

// tokenizerCache caches encodings to avoid repeated initialization overhead
type tokenizerCache struct {
	mu        sync.RWMutex
	encodings map[string]*tiktoken.Tiktoken
}

var cache = &tokenizerCache{encodings: make(map[string]*tiktoken.Tiktoken)}

// GetEncoding returns a Tiktoken encoding for the given encoding name.
func GetEncoding(encodingName string) (*tiktoken.Tiktoken, error) {
	// Check cache first
	cache.mu.RLock()
	if enc, ok := cache.encodings[encodingName]; ok {
		cache.mu.RUnlock()
		return enc, nil
	}
	cache.mu.RUnlock()

	// Initialise encoding
	enc, err := tiktoken.GetEncoding(encodingName)
	if err != nil {
		return nil, fmt.Errorf("failed to get encoding %s: %w", encodingName, err)
	}

	// Cache it
	cache.mu.Lock()
	cache.encodings[encodingName] = enc
	cache.mu.Unlock()

	return enc, nil
}

// EncodingForModel returns the appropriate encoding name for a given model.
func EncodingForModel(model string) (string, error) {
	model = strings.ToLower(model)

	// GPT-4o and newer use o200k_base
	if strings.Contains(model, "gpt-4o") ||
		strings.Contains(model, "gpt-4.1") ||
		strings.Contains(model, "gpt-4.5") ||
		strings.HasPrefix(model, "gpt-4o") ||
		strings.HasPrefix(model, "gpt-4.1") ||
		strings.HasPrefix(model, "gpt-4.5") {
		return EncodingO200kBase, nil
	}

	// GPT-4 family uses cl100k_base
	if strings.HasPrefix(model, "gpt-4") {
		return EncodingCl100kBase, nil
	}

	// GPT-3.5-turbo uses cl100k_base
	if strings.Contains(model, "gpt-3.5-turbo") {
		return EncodingCl100kBase, nil
	}

	// Embedding models use cl100k_base
	if strings.Contains(model, "text-embedding") {
		return EncodingCl100kBase, nil
	}

	// Codex and davinci models use p50k_base
	if strings.Contains(model, "code-") ||
		strings.Contains(model, "davinci") {
		return EncodingP50kBase, nil
	}

	// Default to cl100k_base (most common for modern models)
	return EncodingCl100kBase, nil
}

// CountTokens counts the tokens in a text string using the appropriate encoding.
func CountTokens(text string, model string) (int, error) {
	encodingName, err := EncodingForModel(model)
	if err != nil {
		return 0, err
	}

	enc, err := GetEncoding(encodingName)
	if err != nil {
		return 0, fmt.Errorf("failed to get encoding: %w", err)
	}

	tokens := enc.Encode(text, nil, nil)
	return len(tokens), nil
}