package runtime

import (
	agentllm "agent-native/runtime/llm"
	core "agent-system-core"
	"context"
	"encoding/json"
	"reflect"
	"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")
	}
}

func extractCheckpointDescriptionFindingExamples(desc string) []core.WorkFindingCandidate {
	var out []core.WorkFindingCandidate
	for i := 0; i < len(desc); i++ {
		if !strings.HasPrefix(desc[i:], `{"findingKey"`) {
			continue
		}
		dec := json.NewDecoder(strings.NewReader(desc[i:]))
		var finding core.WorkFindingCandidate
		if err := dec.Decode(&finding); err != nil {
			continue
		}
		out = append(out, finding)
	}
	return out
}

func TestWorkCheckpointDescriptionExamplesValidate(t *testing.T) {
	examples := extractCheckpointDescriptionFindingExamples(checkpointDecl().Description)
	if len(examples) != 4 {
		t.Fatalf("want 4 description examples, got %d", len(examples))
	}
	kinds := map[string]bool{}
	for i, finding := range examples {
		cp := core.WorkCheckpointCandidate{
			CheckpointID: "ex", PhaseID: "reading",
			Findings: []core.WorkFindingCandidate{finding},
		}
		if err := core.ValidateWorkCheckpoint(cp); err != nil {
			t.Fatalf("example %d failed validation: %v", i, err)
		}
		key := finding.Kind
		if finding.Evidence != nil {
			key += "+" + finding.Evidence.Mode
		} else {
			key += "+none"
		}
		kinds[key] = true
	}
	for _, want := range []string{"material_record+none", "material_record+direct_relation", "analysis+difference", "needs_user+unresolved_relation"} {
		if !kinds[want] {
			t.Fatalf("missing allowed pattern example %s in %v", want, kinds)
		}
	}
}

func TestWorkCheckpointDescriptionCounterExamplesRejected(t *testing.T) {
	examples := extractCheckpointDescriptionFindingExamples(checkpointDecl().Description)
	var direct, unresolved *core.WorkFindingCandidate
	for i := range examples {
		f := &examples[i]
		if f.Evidence == nil {
			continue
		}
		switch f.Evidence.Mode {
		case "direct_relation":
			direct = f
		case "unresolved_relation":
			unresolved = f
		}
	}
	if direct == nil || unresolved == nil {
		t.Fatal("missing direct_relation or unresolved_relation example")
	}
	result := 1
	direct.Evidence.Result = &result
	if err := core.ValidateWorkCheckpoint(core.WorkCheckpointCandidate{CheckpointID: "bad-direct", PhaseID: "reading", Findings: []core.WorkFindingCandidate{*direct}}); err == nil {
		t.Fatal("direct_relation with result accepted")
	}
	unresolved.Evidence.Result = &result
	if err := core.ValidateWorkCheckpoint(core.WorkCheckpointCandidate{CheckpointID: "bad-unresolved", PhaseID: "reading", Findings: []core.WorkFindingCandidate{*unresolved}}); err == nil {
		t.Fatal("unresolved_relation with result accepted")
	}
}

func stripJSONDescriptions(v any) any {
	switch x := v.(type) {
	case map[string]any:
		out := make(map[string]any, len(x))
		for k, val := range x {
			if k == "description" {
				continue
			}
			out[k] = stripJSONDescriptions(val)
		}
		return out
	case []any:
		out := make([]any, len(x))
		for i, val := range x {
			out[i] = stripJSONDescriptions(val)
		}
		return out
	default:
		return v
	}
}

func TestWorkCheckpointEvidenceSchemaStructureUnchanged(t *testing.T) {
	const oldSchema = `{
  "type":"object","additionalProperties":false,"required":["checkpointId","phaseId","findings"],
  "properties":{
    "checkpointId":{"type":"string"},
    "phaseId":{"type":"string","enum":["reading","crosscheck","summarizing","draft-ready"]},
    "focus":{"type":"string"},"remaining":{"type":"string"},"nextStep":{"type":"string"},
    "findings":{"type":"array","minItems":1,"maxItems":3,"items":{
      "type":"object","additionalProperties":false,"required":["findingKey","text","kind","sources"],
      "properties":{
        "findingKey":{"type":"string"},"text":{"type":"string"},"impact":{"type":"string"},
        "kind":{"type":"string","enum":["material_record","analysis","provisional","needs_user"]},
        "assumptions":{"type":"string"},"corrects":{"type":"string"},
        "sources":{"type":"array","minItems":1,"items":{"type":"object","additionalProperties":false,
          "required":["sourceId","blockIndex","quote"],"properties":{
            "sourceId":{"type":"string"},"blockIndex":{"type":"integer","minimum":0},
            "quote":{"type":"string","description":"同一段落连续出现的原文短句"}}}},
        "evidence":{"type":"object","additionalProperties":false,"required":["mode","premises"],"properties":{
          "mode":{"type":"string","enum":["direct_relation","difference","unresolved_relation"]},
          "premises":{"type":"array","minItems":1,"maxItems":3,"items":{"type":"object","additionalProperties":false,
            "required":["role","sourceIndex"],"properties":{
              "role":{"type":"string","enum":["total","phase","dependent","direct_relation"]},
              "sourceIndex":{"type":"integer","minimum":0},"value":{"type":"integer","minimum":0,"maximum":1000000}}}},
          "result":{"type":"integer","minimum":0,"maximum":1000000}}}
      }}}
  }
}`
	var oldObj, newObj any
	if err := json.Unmarshal([]byte(oldSchema), &oldObj); err != nil {
		t.Fatal(err)
	}
	if err := json.Unmarshal(checkpointDecl().Parameters, &newObj); err != nil {
		t.Fatal(err)
	}
	if !reflect.DeepEqual(stripJSONDescriptions(oldObj), stripJSONDescriptions(newObj)) {
		t.Fatal("checkpoint schema structure changed beyond description keys")
	}
}

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)
	}
}