package openai

import (
	"strconv"
	"unicode/utf8"
)

// estimateInputTokens returns a conservative approximation of the token count
// for a request payload.  The formula is based on the plan:
//
//	runes/4 + runes/16
//
// where `runes` is the total character count of all serialisable input
// elements (messages, tools, metadata) and structural_overhead is a
// proportional structural overhead.
func estimateInputTokens(input string, metadata map[string]string, tools []any, toolChoice any) int {
	est := estimateFromStrings(input, metadata, tools, toolChoice)
	// Minimum 1 token to avoid division-by-zero issues downstream.
	if est < 1 {
		return 1
	}
	return est
}

// estimateFromStrings produces a conservative token estimate by summing
// string-length proxies plus a small structural overhead.
func estimateFromStrings(input string, metadata map[string]string, tools []any, toolChoice any) int {
	// Count runes in the main text payload.
	totalRunes := 0

	// input is the concatenated prompt/chat text.
	totalRunes += utf8.RuneCountInString(input)

	// metadata values contribute tokens proportional to their size.
	for _, v := range metadata {
		totalRunes += utf8.RuneCountInString(v)
	}

	// tool schemas are typically JSON-heavy; each tool entry adds a base
	// overhead plus the size of its name/description.
	for _, t := range tools {
		totalRunes += 16 // base structural overhead per tool entry
		if fn, ok := t.(map[string]any); ok {
			if name, ok := fn["name"].(string); ok {
				totalRunes += utf8.RuneCountInString(name)
			}
			if parameters, ok := fn["parameters"].(map[string]any); ok {
				totalRunes += utf8.RuneCountInString(toolsSchemaToJSON(parameters))
			}
		}
		if choice, ok := toolChoice.(map[string]any); ok {
			if fn, ok := choice["function"].(map[string]any); ok {
				if name, ok := fn["name"].(string); ok {
					totalRunes += utf8.RuneCountInString(name)
				}
			}
		}
	}

	// Conservative token estimate: each ~4 runes is roughly 1 token.
	tokens := totalRunes / 4
	// Structural overhead: add 1 token per ~16 runes to account for
	// system tokens, delimiter tokens, message boundary tokens, etc.
	tokens += totalRunes / 16
	return tokens
}

// toolsSchemaToJSON is a lightweight JSON stringifier for the tools[].parameters
// object that avoids an import cycle with encoding/json.
func toolsSchemaToJSON(v any) string {
	if m, ok := v.(map[string]any); ok {
		var b stringsBuilder
		b.WriteByte('{')
		first := true
		for k, val := range m {
			if !first {
				b.WriteByte(',')
			}
			first = false
			encodeJSONString(&b, k)
			b.WriteByte(':')
			encodeJSONValue(&b, val)
		}
		b.WriteByte('}')
		return b.String()
	}
	return "{}"
}

// stringsBuilder is a minimal []byte-based string builder to avoid import cycles.
type stringsBuilder struct {
	b []byte
}

func (s *stringsBuilder) WriteByte(b byte) {
	s.b = append(s.b, b)
}

func (s *stringsBuilder) WriteString(str string) {
	s.b = append(s.b, str...)
}

func (s *stringsBuilder) WriteRune(r rune) {
	buf := make([]byte, utf8.UTFMax)
	n := utf8.EncodeRune(buf, r)
	s.b = append(s.b, buf[:n]...)
}

func (s *stringsBuilder) Len() int {
	return len(s.b)
}

func (s *stringsBuilder) String() string {
	return string(s.b)
}

func encodeJSONString(b *stringsBuilder, v string) {
	b.WriteByte('"')
	for i := 0; i < len(v); i++ {
		c := v[i]
		switch c {
		case '"', '\\':
			b.WriteByte('\\')
			b.WriteByte(c)
		default:
			b.WriteByte(c)
		}
	}
	b.WriteByte('"')
}

func encodeJSONValue(b *stringsBuilder, v any) {
	switch val := v.(type) {
	case string:
		encodeJSONString(b, val)
	case bool:
		if val {
			b.WriteString("true")
		} else {
			b.WriteString("false")
		}
	case float64:
		b.WriteString(ftoa(val))
	case int:
		b.WriteString(itoa(val))
	case int64:
		b.WriteString(itoa64(val))
	case map[string]any:
		encodeJSONObject(b, val)
	case []any:
		encodeJSONArray(b, val)
	case nil:
		b.WriteString("null")
	default:
		b.WriteString("null")
	}
}

func encodeJSONObject(b *stringsBuilder, m map[string]any) {
	b.WriteByte('{')
	first := true
	for k, v := range m {
		if !first {
			b.WriteByte(',')
		}
		first = false
		encodeJSONString(b, k)
		b.WriteByte(':')
		encodeJSONValue(b, v)
	}
	b.WriteByte('}')
}

func encodeJSONArray(b *stringsBuilder, items []any) {
	b.WriteByte('[')
	for i, item := range items {
		if i > 0 {
			b.WriteByte(',')
		}
		encodeJSONValue(b, item)
	}
	b.WriteByte(']')
}

func ftoa(v float64) string {
	return strconv.FormatFloat(v, 'f', -1, 64)
}

// lightweight integer/float formatters to avoid fmt import.
func itoa(v int) string {
	if v == 0 {
		return "0"
	}
	negative := v < 0
	if negative {
		v = -v
	}
	var buf [20]byte
	i := len(buf)
	for v > 0 {
		i--
		buf[i] = byte('0' + v%10)
		v /= 10
	}
	if negative {
		i--
		buf[i] = '-'
	}
	return string(buf[i:])
}

func itoa64(v int64) string {
	if v == 0 {
		return "0"
	}
	negative := v < 0
	if negative {
		v = -v
	}
	var buf [20]byte
	i := len(buf)
	for v > 0 {
		i--
		buf[i] = byte('0' + v%10)
		v /= 10
	}
	if negative {
		i--
		buf[i] = '-'
	}
	return string(buf[i:])
}

func zeroPadLeft(pad int, s string) string {
	for len(s) < pad {
		s = "0" + s
	}
	return s
}

// classifyContext returns "normal" or "long" based on the token estimate
// and the server's threshold.
func classifyContext(tokens int, threshold int) string {
	if tokens >= threshold {
		return "long"
	}
	return "normal"
}