package runtime

import (
	"context"
	"testing"

	agentllm "git-ssh.u1680832.nyat.app/weichen/Agent-Native.git/runtime/llm"
	core "git-ssh.u1680832.nyat.app/weichen/Agent-System-Core.git"
)

// ---- N3 WP6：计费拆分与重试收口（backend 侧）----

// retryAttemptsOf 映射：设置 0=不重试→-1；N→N；默认 2→2。
func TestRetryAttemptsMapping(t *testing.T) {
	if got := retryAttemptsOf(core.EffectiveSettings{}); got != 2 {
		t.Fatalf("默认应 2，得到 %d", got)
	}
	if got := retryAttemptsOf(core.EffectiveSettings{"runtime": map[string]any{"llmRetries": float64(0)}}); got != -1 {
		t.Fatalf("设置 0 应映射为关闭（-1），得到 %d", got)
	}
	if got := retryAttemptsOf(core.EffectiveSettings{"runtime": map[string]any{"llmRetries": float64(5)}}); got != 5 {
		t.Fatalf("设置 5 应映射为 5，得到 %d", got)
	}
}

// 直连 run 的用量拆分：message_end 到达后 Last* 覆盖、cache 累计入账。
func TestDirectRunUsageSplitOnRecord(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 1)
		ch <- agentllm.Event{Kind: agentllm.KindMessageEnd,
			Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "x", StopReason: agentllm.StopEnd},
			Usage:   agentllm.Usage{PromptTokens: 120, CompletionTokens: 12, TotalTokens: 132, CacheReadTokens: 80, CacheWriteTokens: 6},
		}
		close(ch)
		return ch
	}
	r := &LLMRunner{Stream: stream}
	if err := r.Execute(rc); err != nil {
		t.Fatalf("Execute: %v", err)
	}
	u := rc.usageTotal
	if u.PromptTokens != 120 || u.CompletionTokens != 12 ||
		u.LastPromptTokens != 120 || u.LastCompletionTokens != 12 ||
		u.CacheReadTokens != 80 || u.CacheWriteTokens != 6 {
		t.Fatalf("用量拆分未入账：%+v", u)
	}
}

// 子代理记账：subRun 的 message_end 累计 token 并经 AddUsage 入账。
func TestSubRunAccountsTokens(t *testing.T) {
	rc := &fakeRunCtx{ctx: context.Background(), prof: core.AgentProfile{Model: core.ModelRoute{Chat: "test-model"}}}
	stream := func(ctx context.Context, model string, msgs []agentllm.Message, _ []agentllm.ToolDecl, _ agentllm.StreamOptions) <-chan agentllm.Event {
		ch := make(chan agentllm.Event, 1)
		ch <- agentllm.Event{Kind: agentllm.KindMessageEnd,
			Message: agentllm.Message{Role: agentllm.RoleAssistant, Content: "子结果", StopReason: agentllm.StopEnd},
			Usage:   agentllm.Usage{PromptTokens: 66, CompletionTokens: 6, TotalTokens: 72, CacheReadTokens: 30},
		}
		close(ch)
		return ch
	}
	r := &SubLLMRunner{Stream: stream}
	if err := r.Execute(rc); err != nil {
		t.Fatalf("Execute: %v", err)
	}
	u := rc.usageTotal
	if u.PromptTokens != 66 || u.CompletionTokens != 6 ||
		u.LastPromptTokens != 66 || u.LastCompletionTokens != 6 || u.CacheReadTokens != 30 {
		t.Fatalf("子代理用量未入账：%+v", u)
	}
}