package harness

import (
	"context"
	"testing"
)

func TestEstimateTokens(t *testing.T) {
	if got := EstimateTokens(""); got != 0 {
		t.Errorf("空串应 0, got %d", got)
	}
	// 8 个 ASCII → (8+3)/4 = 2
	if got := EstimateTokens("abcdefgh"); got != 2 {
		t.Errorf("ASCII 估算 want 2, got %d", got)
	}
	// 4 个中文 → 4
	if got := EstimateTokens("你好世界"); got != 4 {
		t.Errorf("中文估算 want 4, got %d", got)
	}
}

func TestBudget_ExceedAndSnapshot(t *testing.T) {
	b := NewBudget(10)
	b.AddPrompt("你好世界") // 4
	b.AddComplete("你好") // 2 → total 6
	if b.Exceeded() {
		t.Fatal("6/10 不应触顶")
	}
	b.AddComplete("一二三四五") // +5 → 11
	if !b.Exceeded() {
		t.Fatal("11/10 应触顶")
	}
	p, c, total := b.Snapshot()
	if p != 4 || c != 7 || total != 11 {
		t.Errorf("快照 want 4/7/11, got %d/%d/%d", p, c, total)
	}
}

func TestBudget_Unlimited(t *testing.T) {
	b := NewBudget(0) // 不限额
	b.AddComplete("非常非常非常长的一段中文内容反复堆叠占用大量预算额度")
	if b.Exceeded() {
		t.Error("max≤0 不应触顶")
	}
}

func TestBudget_NilSafe(t *testing.T) {
	var b *Budget
	b.AddPrompt("x") // 不应 panic
	if b.Exceeded() {
		t.Error("nil 预算视为不限额")
	}
	if _, _, total := b.Snapshot(); total != 0 {
		t.Error("nil 快照应为 0")
	}
}

func TestBudget_Context(t *testing.T) {
	if BudgetFrom(context.Background()) != nil {
		t.Error("无预算的 ctx 应返回 nil")
	}
	b := NewBudget(100)
	ctx := WithBudget(context.Background(), b)
	if BudgetFrom(ctx) != b {
		t.Error("应取回同一预算实例")
	}
}