package harness

import (
	"context"
	"sync"
	"unicode"
)

// EstimateTokens 粗估一段文本的 token 数（无需依赖模型/分词器，作预算护栏足够）：
// CJK 字符约 1 token/字；其余（拉丁/数字/空白/标点）约 1 token/4 字符。偏保守（宁高勿低）。
func EstimateTokens(s string) int {
	cjk, other := 0, 0
	for _, r := range s {
		if unicode.Is(unicode.Han, r) || unicode.Is(unicode.Hiragana, r) ||
			unicode.Is(unicode.Katakana, r) || unicode.Is(unicode.Hangul, r) {
			cjk++
		} else {
			other++
		}
	}
	return cjk + (other+3)/4 // 向上取整
}

// Budget 是单任务的 token 预算计量与封顶（并发安全：流式回调与节点循环可能并发计量）。
// max ≤ 0 表示不限额（仅计量、不中止）。
type Budget struct {
	mu     sync.Mutex
	max    int
	prompt int
	comp   int
}

// NewBudget 建一个上限为 max 的任务预算（max≤0 → 不限额，纯计量）。
func NewBudget(max int) *Budget { return &Budget{max: max} }

// AddPrompt / AddComplete 累计输入/输出 token（按文本估算）。
func (b *Budget) AddPrompt(text string)   { b.add(EstimateTokens(text), 0) }
func (b *Budget) AddComplete(text string) { b.add(0, EstimateTokens(text)) }

func (b *Budget) add(p, c int) {
	if b == nil {
		return
	}
	b.mu.Lock()
	b.prompt += p
	b.comp += c
	b.mu.Unlock()
}

// Exceeded 报告是否已触顶（max≤0 恒为 false）。
func (b *Budget) Exceeded() bool {
	if b == nil {
		return false
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.max > 0 && b.prompt+b.comp > b.max
}

// Snapshot 返回当前用量快照（prompt / completion / total）。
func (b *Budget) Snapshot() (prompt, comp, total int) {
	if b == nil {
		return 0, 0, 0
	}
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.prompt, b.comp, b.prompt + b.comp
}

// budgetKey 是 context 里承载任务预算的私有键（避免跨包碰撞）。
type budgetKey struct{}

// WithBudget 把任务预算挂到 context，沿图执行透传（免改各节点函数签名）。
func WithBudget(ctx context.Context, b *Budget) context.Context {
	return context.WithValue(ctx, budgetKey{}, b)
}

// BudgetFrom 取出 context 里的任务预算（无则返回 nil，调用方按"不限额"处理）。
func BudgetFrom(ctx context.Context) *Budget {
	b, _ := ctx.Value(budgetKey{}).(*Budget)
	return b
}