package stream

import (
	"context"
	"strings"
	"sync"
	"time"
	"unicode"

	"aichat/message"
)

type Stats struct {
	PromptTokens               int     `json:"prompt_tokens"`
	CompletionTokens           int     `json:"completion_tokens"`
	ToolPromptTokens           int     `json:"tool_prompt_tokens"`
	ToolCompletionTokens       int     `json:"tool_completion_tokens"`
	TotalTokens                int     `json:"total_tokens"`
	CompletionTokensPerSec     float64 `json:"completion_tokens_per_sec"`
	PeakCompletionTokensPerSec float64 `json:"peak_completion_tokens_per_sec"`
	Estimated                  bool    `json:"estimated"`
}

type Tracker struct {
	mu                   sync.Mutex
	promptTokens         int
	completionTokens     int
	toolPromptTokens     int
	toolCompletionTokens int
}

type tokenUsageContextKey struct{}

func NewTracker() *Tracker {
	return &Tracker{}
}

func ContextWithTracker(ctx context.Context, tracker *Tracker) context.Context {
	if tracker == nil {
		return ctx
	}
	return context.WithValue(ctx, tokenUsageContextKey{}, tracker)
}

func TrackerFromContext(ctx context.Context) *Tracker {
	tracker, _ := ctx.Value(tokenUsageContextKey{}).(*Tracker)
	return tracker
}

func (t *Tracker) AddTool(promptTokens, completionTokens int) {
	if t == nil {
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	t.toolPromptTokens += promptTokens
	t.toolCompletionTokens += completionTokens
}

func (t *Tracker) SetModel(promptTokens, completionTokens int) {
	if t == nil {
		return
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	t.promptTokens = promptTokens
	t.completionTokens = completionTokens
}

func (t *Tracker) Snapshot(tokensPerSecond, peakTokensPerSecond float64) Stats {
	if t == nil {
		return Stats{Estimated: true}
	}
	t.mu.Lock()
	defer t.mu.Unlock()
	total := t.promptTokens + t.completionTokens + t.toolPromptTokens + t.toolCompletionTokens
	return Stats{
		PromptTokens:               t.promptTokens,
		CompletionTokens:           t.completionTokens,
		ToolPromptTokens:           t.toolPromptTokens,
		ToolCompletionTokens:       t.toolCompletionTokens,
		TotalTokens:                total,
		CompletionTokensPerSec:     tokensPerSecond,
		PeakCompletionTokensPerSec: peakTokensPerSecond,
		Estimated:                  true,
	}
}

func EstimateChatMessagesTokens(messages []message.ChatMessage) int {
	total := 0
	for _, msg := range messages {
		total += EstimateTokenCount(msg.Role) + EstimateTokenCount(msg.Content) + 4
		if msg.ImageURL != "" || msg.ImageURLAlias != "" {
			total += 85
		}
	}
	return total
}

func EstimateTokenCount(text string) int {
	text = strings.TrimSpace(text)
	if text == "" {
		return 0
	}
	tokens := 0
	asciiRunes := 0
	flushASCII := func() {
		if asciiRunes > 0 {
			tokens += (asciiRunes + 3) / 4
			asciiRunes = 0
		}
	}
	for _, r := range text {
		if unicode.IsSpace(r) {
			flushASCII()
			continue
		}
		if r <= unicode.MaxASCII {
			asciiRunes++
			continue
		}
		flushASCII()
		tokens++
	}
	flushASCII()
	if tokens == 0 {
		return 1
	}
	return tokens
}

func TokensPerSecond(tokens int, start time.Time) float64 {
	elapsed := time.Since(start).Seconds()
	if tokens <= 0 || elapsed <= 0 {
		return 0
	}
	return float64(tokens) / elapsed
}