package runtime

import (
	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
	"context"
	"encoding/json"
	"strings"
	"testing"
)

func TestMaterialFinalReserveAfterTwoCheckpoints(t *testing.T) {
	limit := 100000
	prof := subProfile()
	prof.Budget.MaxTokens = &limit
	rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID, task: "核对项目材料"}}}
	turn := func(body string, call ...agentllm.ToolCall) []agentllm.Event {
		events := textDeltaTurn(body, call...)
		events[len(events)-1].Usage = agentllm.Usage{PromptTokens: 4000, CompletionTokens: 1000, TotalTokens: 5000}
		return events
	}
	base := scriptStream([][]agentllm.Event{
		turn("", toolCall(workCheckpointTool, checkpointJSON())),
		turn("", toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp2"`, 1))),
		turn("整理完成"),
	}, nil)
	calls, finalOnly, identityStable := 0, false, false
	firstSystem, firstUser := "", ""
	r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		calls++
		if calls == 1 {
			for _, msg := range msgs {
				if msg.Role == agentllm.RoleSystem {
					firstSystem = msg.Content
				}
				if msg.Role == agentllm.RoleUser {
					firstUser = msg.Content
					break
				}
			}
		}
		if calls == 3 {
			identityStable = firstSystem != "" && firstUser != "" && len(msgs) > 0 && msgs[0].Role == agentllm.RoleSystem && msgs[0].Content == firstSystem
			for _, msg := range msgs {
				if msg.Role == agentllm.RoleUser {
					identityStable = identityStable && msg.Content == firstUser
					break
				}
			}
			finalOnly = true
			for _, d := range decls {
				if d.Name == workCheckpointTool {
					finalOnly = false
				}
			}
			reminder := false
			for _, msg := range msgs {
				reminder = reminder || (msg.Role == agentllm.RoleSystem && strings.Contains(msg.Content, "不再发布检查点"))
			}
			finalOnly = finalOnly && reminder
		}
		return base(ctx, model, msgs, decls, opts)
	}}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if calls != 3 || len(rc.checkpoints) != 2 || !finalOnly || !identityStable {
		t.Fatalf("final-only transition or evaluator identity missing: calls=%d checkpoints=%d finalOnly=%v identityStable=%v", calls, len(rc.checkpoints), finalOnly, identityStable)
	}
}

func TestMaterialFinalReserveRejectsHiddenCheckpointWithoutAnotherProviderCall(t *testing.T) {
	limit := 100000
	prof := subProfile()
	prof.Budget.MaxTokens = &limit
	rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}}
	turn := func(call agentllm.ToolCall) []agentllm.Event {
		events := textDeltaTurn("", call)
		events[len(events)-1].Usage = agentllm.Usage{PromptTokens: 4000, CompletionTokens: 1000, TotalTokens: 5000}
		return events
	}
	base := scriptStream([][]agentllm.Event{
		turn(toolCall(workCheckpointTool, checkpointJSON())),
		turn(toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp2"`, 1))),
		turn(toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp3"`, 1))),
		textDeltaTurn("unexpected extra response"),
	}, nil)
	providerCalls := 0
	r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		providerCalls++
		return base(ctx, model, msgs, decls, opts)
	}}
	_ = r.Execute(rc)
	if providerCalls != 3 || len(rc.checkpoints) != 2 {
		t.Fatalf("hidden checkpoint caused a publication or another provider call: provider=%d saved=%d", providerCalls, len(rc.checkpoints))
	}
}

func TestMaterialFinalReserveCountsDistinctPersistedCheckpointIDs(t *testing.T) {
	limit := 100000
	prof := subProfile()
	prof.Budget.MaxTokens = &limit
	rc := &cappedIdempotentCheckpointRC{&idempotentCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}}}
	turn := func(body string, call ...agentllm.ToolCall) []agentllm.Event {
		events := textDeltaTurn(body, call...)
		events[len(events)-1].Usage = agentllm.Usage{PromptTokens: 4000, CompletionTokens: 1000, TotalTokens: 5000}
		return events
	}
	base := scriptStream([][]agentllm.Event{
		turn("", toolCall(workCheckpointTool, checkpointJSON())),
		turn("", toolCall(workCheckpointTool, checkpointJSON())),
		turn("", toolCall(workCheckpointTool, strings.Replace(checkpointJSON(), `"cp1"`, `"cp2"`, 1))),
		turn("整理完成"),
	}, nil)
	providerCalls, thirdCanPublish, fourthFinalOnly := 0, false, false
	r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		providerCalls++
		canPublish := false
		for _, d := range decls {
			canPublish = canPublish || d.Name == workCheckpointTool
		}
		if providerCalls == 3 {
			thirdCanPublish = canPublish
		}
		if providerCalls == 4 {
			fourthFinalOnly = !canPublish
		}
		return base(ctx, model, msgs, decls, opts)
	}}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if providerCalls != 4 || len(rc.checkpoints) != 2 || !thirdCanPublish || !fourthFinalOnly {
		t.Fatalf("idempotent replay changed final-only transition: calls=%d saved=%d thirdCanPublish=%v fourthFinalOnly=%v", providerCalls, len(rc.checkpoints), thirdCanPublish, fourthFinalOnly)
	}
}

func TestMaterialFinalReserveAnswersAfterOneCheckpointWhenBudgetNarrows(t *testing.T) {
	limit := 30000
	prof := subProfile()
	prof.Budget.MaxTokens = &limit
	rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}}
	first := textDeltaTurn("", toolCall(workCheckpointTool, checkpointJSON()))
	first[len(first)-1].Usage = agentllm.Usage{PromptTokens: 8000, CompletionTokens: 2000, TotalTokens: 10000}
	base := scriptStream([][]agentllm.Event{first, textDeltaTurn("目前只核对出一项，其他内容尚待核对")}, nil)
	providerCalls, finalOnly := 0, false
	r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		providerCalls++
		if providerCalls == 2 {
			finalOnly = true
			for _, d := range decls {
				finalOnly = finalOnly && d.Name != workCheckpointTool
			}
		}
		return base(ctx, model, msgs, decls, opts)
	}}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if providerCalls != 2 || len(rc.checkpoints) != 1 || !finalOnly {
		t.Fatalf("partial final answer was not requested: calls=%d saved=%d finalOnly=%v", providerCalls, len(rc.checkpoints), finalOnly)
	}
}

func TestMaterialFinalReserveStopsBeforeUnaffordableRequest(t *testing.T) {
	for _, tc := range []struct {
		name  string
		usage agentllm.Usage
	}{
		{"reported usage exceeds final prompt allowance", agentllm.Usage{PromptTokens: 25000, CompletionTokens: 2000, TotalTokens: 27000}},
		{"missing usage", agentllm.Usage{}},
	} {
		t.Run(tc.name, func(t *testing.T) {
			limit := 30000
			prof := subProfile()
			prof.Budget.MaxTokens = &limit
			rc := &cappedCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: materialSubagentID}}}
			first := textDeltaTurn("", toolCall(workCheckpointTool, checkpointJSON()))
			first[len(first)-1].Usage = tc.usage
			base := scriptStream([][]agentllm.Event{first, textDeltaTurn("unexpected extra response")}, nil)
			providerCalls := 0
			r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
				providerCalls++
				return base(ctx, model, msgs, decls, opts)
			}}
			if err := r.Execute(rc); err == nil {
				t.Fatal("expected honest budget/usage failure")
			}
			if providerCalls != 1 || len(rc.checkpoints) != 1 {
				t.Fatalf("unaffordable request made or saved checkpoint lost: calls=%d saved=%d", providerCalls, len(rc.checkpoints))
			}
		})
	}
}

func TestMaterialFinalReserveDoesNotChangeOrdinarySubagents(t *testing.T) {
	limit := 10000
	prof := subProfile()
	prof.Budget.MaxTokens = &limit
	rc := &subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: prof}, id: "agent.sub.explore", tools: subTools()}
	first := textDeltaTurn("", toolCall("knowledge_search", `{"query":"x"}`))
	first[len(first)-1].Usage = agentllm.Usage{PromptTokens: 5000, CompletionTokens: 1000, TotalTokens: 6000}
	base := scriptStream([][]agentllm.Event{first, textDeltaTurn("普通子代理已完成")}, nil)
	providerCalls, reminder := 0, false
	r := &SubLLMRunner{Stream: func(ctx context.Context, model string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		providerCalls++
		for _, msg := range msgs {
			reminder = reminder || strings.Contains(msg.Content, "不再发布检查点")
		}
		return base(ctx, model, msgs, decls, opts)
	}}
	if err := r.Execute(rc); err != nil {
		t.Fatal(err)
	}
	if providerCalls != 2 || reminder {
		t.Fatalf("ordinary subagent changed: provider calls=%d material reminder=%v", providerCalls, reminder)
	}
}

type checkpointRC struct {
	subRC
	checkpoints []core.WorkCheckpointCandidate
}

type receiptCheckpointRC struct{ *checkpointRC }

func (c *receiptCheckpointRC) PublishWorkCheckpoint(v core.WorkCheckpointCandidate) (string, error) {
	c.checkpoints = append(c.checkpoints, v)
	return "集合归属不唯一，已转待确认", nil
}

func TestWorkCheckpointEvidenceReceiptReachesExistingToolResult(t *testing.T) {
	rc := &receiptCheckpointRC{&checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()}, id: materialSubagentID}}}
	result := publishCheckpoint(rc, checkpointJSON())
	if result.IsError || !strings.Contains(result.Content, "待确认") || len(rc.checkpoints) != 1 {
		t.Fatalf("normalization receipt absent from existing tool result: %+v", result)
	}
}

func TestWorkCheckpointEvidenceToolSchemaIsStrictAndBounded(t *testing.T) {
	var schema map[string]any
	if err := json.Unmarshal(checkpointDecl().Parameters, &schema); err != nil {
		t.Fatal("invalid checkpoint tool schema", err)
	}
	root := schema["properties"].(map[string]any)
	finding := root["findings"].(map[string]any)["items"].(map[string]any)
	if finding["additionalProperties"] != false {
		t.Fatal("finding admits unknown fields")
	}
	evidence := finding["properties"].(map[string]any)["evidence"].(map[string]any)
	if evidence["additionalProperties"] != false {
		t.Fatal("evidence admits unknown fields")
	}
	premises := evidence["properties"].(map[string]any)["premises"].(map[string]any)
	if premises["maxItems"] != float64(3) || premises["items"].(map[string]any)["additionalProperties"] != false {
		t.Fatal("premises are unbounded or admit unknown fields")
	}
}

type cappedCheckpointRC struct{ *checkpointRC }

type idempotentCheckpointRC struct{ *checkpointRC }

func (c *idempotentCheckpointRC) PublishWorkCheckpoint(v core.WorkCheckpointCandidate) (string, error) {
	for _, saved := range c.checkpoints {
		if saved.CheckpointID == v.CheckpointID {
			return "", nil
		}
	}
	return c.checkpointRC.PublishWorkCheckpoint(v)
}

type cappedIdempotentCheckpointRC struct{ *idempotentCheckpointRC }

func (c *cappedIdempotentCheckpointRC) Settings() core.EffectiveSettings {
	return (&cappedCheckpointRC{c.checkpointRC}).Settings()
}

func (c *cappedCheckpointRC) Settings() core.EffectiveSettings {
	eff := c.checkpointRC.Settings()
	eff["models"].(map[string]any)["chat"].(map[string]any)["maxOutputTokens"] = float64(4000)
	return eff
}

func (c *checkpointRC) PublishWorkCheckpoint(v core.WorkCheckpointCandidate) (string, error) {
	c.checkpoints = append(c.checkpoints, v)
	return "", nil
}
func checkpointJSON() string {
	return `{"checkpointId":"cp1","phaseId":"reading","focus":"核对决定","findings":[{"findingKey":"date","text":"日期有差异","kind":"material_record","sources":[{"sourceId":"s1","blockIndex":0,"quote":"上线六月"}]}]}`
}
func TestWorkCheckpointUsesLocalSinkAndNoExtraModelCall(t *testing.T) {
	rc := &checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()}, id: materialSubagentID}}
	calls := 0
	base := scriptStream([][]agentllm.Event{textDeltaTurn("", toolCall("publish_work_checkpoint", checkpointJSON())), textDeltaTurn("完成")}, nil)
	r := &SubLLMRunner{Stream: func(ctx context.Context, m string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		calls++
		found := false
		for _, d := range decls {
			if d.Name == "publish_work_checkpoint" {
				found = true
			}
		}
		if !found {
			t.Error("material checkpoint tool missing")
		}
		return base(ctx, m, msgs, decls, opts)
	}}
	if e := r.Execute(rc); e != nil {
		t.Fatal(e)
	}
	if len(rc.checkpoints) != 1 || rc.toolRuns != 0 || calls != 2 || rc.usageTotal.ToolCalls != 1 {
		t.Fatal("checkpoint wasn't local / extra summary call", len(rc.checkpoints), rc.toolRuns, calls)
	}
}
func TestWorkCheckpointPolicyAndStrictArguments(t *testing.T) {
	rc := &checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()}, id: materialSubagentID}}
	for _, policy := range []core.ToolPolicy{{Deny: []string{"publish_work_checkpoint"}}, {Confirm: []string{"publish_work_checkpoint"}}} {
		rc.prof.Tools = policy
		for _, d := range subagentDecls(rc, materialSubagentSpec()) {
			if d.Name == "publish_work_checkpoint" {
				t.Fatal("policy bypass")
			}
		}
	}
	rc.prof.Tools = core.ToolPolicy{}
	st := &subRun{rc: rc, localCheckpoint: checkpointEnabled(rc)}
	res := st.runTool("publish_work_checkpoint", `{"checkpointId":"cp1","taskId":"forged"}`)
	if !res.IsError || len(rc.checkpoints) != 0 {
		t.Fatal("foreign scope accepted", res)
	}
}

func TestWorkCheckpointFailureDoesNotFailMainTask(t *testing.T) {
	rc := &checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()}, id: materialSubagentID}}
	r := &SubLLMRunner{Stream: scriptStream([][]agentllm.Event{textDeltaTurn("", toolCall(workCheckpointTool, `{"checkpointId":"bad","taskId":"forged"}`)), textDeltaTurn("保留最终结果")}, nil)}
	if e := r.Execute(rc); e != nil {
		t.Fatal(e)
	}
	if len(rc.checkpoints) != 0 {
		t.Fatal("bad candidate published")
	}
	found := false
	for _, p := range rc.assistant {
		if p.Text == "保留最终结果" {
			found = true
		}
	}
	if !found {
		t.Fatal("main result lost")
	}
}
func TestWorkCheckpointMaterialRespectsProfileTurnBudget(t *testing.T) {
	one := 1
	rc := &checkpointRC{subRC: subRC{fakeRunCtx: fakeRunCtx{ctx: context.Background(), prof: subProfile()}, id: materialSubagentID}}
	rc.prof.Budget.MaxTurns = &one
	calls := 0
	base := scriptStream([][]agentllm.Event{textDeltaTurn("", toolCall(workCheckpointTool, checkpointJSON())), textDeltaTurn("must not run")}, nil)
	r := &SubLLMRunner{Stream: func(ctx context.Context, m string, msgs []agentllm.Message, decls []agentllm.ToolDecl, opts agentllm.StreamOptions) <-chan agentllm.Event {
		calls++
		return base(ctx, m, msgs, decls, opts)
	}}
	_ = r.Execute(rc)
	if calls != 1 {
		t.Fatal("child ignored profile budget", calls)
	}
}