package chunking

import (
	"unicode"
)

// TokenEstimator provides a deterministic token estimation suitable for prompt budgeting.
// The estimator is approximate but stable, isolated, and replaceable.
type TokenEstimator interface {
	EstimateTokens(text string) int
}

// SimpleTokenEstimator provides a basic deterministic token estimation.
// This uses a simple heuristic based on word count and punctuation.
type SimpleTokenEstimator struct{}

// NewSimpleTokenEstimator creates a new simple token estimator.
func NewSimpleTokenEstimator() *SimpleTokenEstimator {
	return &SimpleTokenEstimator{}
}

// EstimateTokens provides a rough estimate of the number of tokens in the given text.
// This implementation uses a simple heuristic: count words and punctuation as tokens.
// The estimate is deterministic and stable for the same input text.
func (e *SimpleTokenEstimator) EstimateTokens(text string) int {
	words, punctuation := e.CountComponents(text)
	return words + punctuation/2
}

// CountComponents returns unrounded counts for budgets composed from fragments.
func (e *SimpleTokenEstimator) CountComponents(text string) (words, punctuation int) {
	inWord := false
	for _, r := range text {
		nonspace := !unicode.IsSpace(r)
		if nonspace && !inWord {
			words++
		}
		inWord = nonspace
		if unicode.IsPunct(r) && r != '\'' && r != '-' && r != '_' {
			punctuation++
		}
	}
	return words, punctuation
}

// ConstTokenEstimator returns a constant token count for testing purposes.
type ConstTokenEstimator struct {
	Tokens int
}

// EstimateTokens returns the configured constant token count.
func (e *ConstTokenEstimator) EstimateTokens(text string) int {
	return e.Tokens
}