package agent

import (
	"encoding/json"
	"math/bits"
	"sort"
	"strings"
	"sync"
	"unicode"

	"github.com/pkoukk/tiktoken-go"

	"github.com/tidbcloud/AgentTeam/internal/apr/tokenencoding"
)

var tokenSourceTokenizers sync.Map

// estimateTokenSourceContent keeps only a numeric proxy in Process state. The
// raw value is discarded immediately after counting.
func estimateTokenSourceContent(value any, runtime string) uint64 {
	switch typed := value.(type) {
	case nil:
		return 0
	case string:
		return estimateTokenSourceText(typed, runtime)
	case []byte:
		return estimateTokenSourceText(string(typed), runtime)
	case json.RawMessage:
		return estimateTokenSourceText(string(typed), runtime)
	default:
		encoded, err := json.Marshal(typed)
		if err != nil {
			return 0
		}
		return estimateTokenSourceText(string(encoded), runtime)
	}
}

// estimateTokenSourceText uses the closest bundled offline tokenizer for the
// runtime family. Provider totals remain authoritative; this local count is
// only an allocation weight and falls back to a deterministic heuristic.
func estimateTokenSourceText(value string, runtime string) uint64 {
	if value == "" {
		return 0
	}
	encoding := "o200k_base"
	if strings.EqualFold(strings.TrimSpace(runtime), "claude") {
		encoding = "cl100k_base"
	}
	if tokenizer := tokenSourceTokenizer(encoding); tokenizer != nil {
		return uint64(len(tokenizer.EncodeOrdinary(value)))
	}
	return estimateTokenSourceTextFallback(value)
}

func tokenSourceTokenizer(encoding string) *tiktoken.Tiktoken {
	if cached, ok := tokenSourceTokenizers.Load(encoding); ok {
		return cached.(*tiktoken.Tiktoken)
	}
	tokenizer, err := tokenencoding.GetEncoding(encoding)
	if err != nil {
		return nil
	}
	actual, _ := tokenSourceTokenizers.LoadOrStore(encoding, tokenizer)
	return actual.(*tiktoken.Tiktoken)
}

func estimateTokenSourceTextFallback(value string) uint64 {
	var tokens uint64
	inWord := false
	for _, r := range value {
		switch {
		case unicode.IsSpace(r):
			inWord = false
		case r < 0x80 && (unicode.IsLetter(r) || unicode.IsDigit(r) || r == '_'):
			if !inWord {
				tokens++
				inWord = true
			}
		default:
			tokens++
			inWord = false
		}
	}
	return tokens
}

func allocateTokenSourceObservations(
	base TokenSourceObservation,
	consumption tokenSourceConsumption,
) []TokenSourceObservation {
	if consumption.decision == "exact" {
		base.Source = consumption.source
		base.CallCount = 1
		return []TokenSourceObservation{base}
	}
	if len(consumption.candidates) == 0 {
		base.Source = consumption.source
		if base.Source.Kind == tokenSourceKindTool || base.Source.Kind == tokenSourceKindCommand {
			base.CallCount = 1
		}
		return []TokenSourceObservation{base}
	}

	// Cached input is an opaque provider subset: the runtime does not identify
	// which cached fragments came from which tool result. Only fresh input is
	// eligible for estimated tool allocation; the cached subset stays context.
	allocatableInput := base.InputTokens
	if base.CachedInputTokens < allocatableInput {
		allocatableInput -= base.CachedInputTokens
	} else {
		allocatableInput = 0
	}
	inputShares := largestRemainderTokenShares(allocatableInput, consumption.candidates)
	observations := make([]TokenSourceObservation, 0, len(consumption.candidates)+1)
	var allocatedInput uint64
	for index, candidate := range consumption.candidates {
		observation := base
		observation.Source = candidate.source
		observation.InputTokens = inputShares[index]
		observation.CachedInputTokens = 0
		observation.OutputTokens = 0
		observation.CallCount = 1
		allocatedInput += inputShares[index]
		observations = append(observations, observation)
	}

	context := base
	context.Source = contextTokenSource()
	context.InputTokens = base.InputTokens - allocatedInput
	context.CallCount = 0
	observations = append(observations, context)
	return observations
}

type tokenSourceRemainder struct {
	index     int
	remainder uint64
}

func largestRemainderTokenShares(
	total uint64,
	candidates []tokenSourceCandidate,
) []uint64 {
	shares := make([]uint64, len(candidates))
	var weightTotal uint64
	for _, candidate := range candidates {
		next, carry := bits.Add64(weightTotal, candidate.outputTokens, 0)
		if carry != 0 {
			// Real estimates are bounded by in-memory string size, but fail
			// safely if a synthetic or future producer exceeds uint64.
			weightTotal = ^uint64(0)
			break
		}
		weightTotal = next
	}
	if total == 0 || weightTotal == 0 {
		return shares
	}
	if weightTotal <= total {
		for index, candidate := range candidates {
			shares[index] = candidate.outputTokens
		}
		return shares
	}

	remainders := make([]tokenSourceRemainder, len(candidates))
	var allocated uint64
	for index, candidate := range candidates {
		high, low := bits.Mul64(total, candidate.outputTokens)
		shares[index], low = bits.Div64(high, low, weightTotal)
		allocated += shares[index]
		remainders[index] = tokenSourceRemainder{
			index:     index,
			remainder: low,
		}
	}
	sort.SliceStable(remainders, func(i, j int) bool {
		return remainders[i].remainder > remainders[j].remainder
	})
	for index := uint64(0); index < total-allocated; index++ {
		shares[remainders[index].index]++
	}
	return shares
}