package agent

import (
	"context"
	"fmt"
	"regexp"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.gopher.ink/mirrors2/AutoZero/llm"
	"git.gopher.ink/mirrors2/AutoZero/tools"
)

// subAgentTester 建一个父 agent 并挂上 agent 工具，返回父 agent 与工具本身
// （用例常要往父级上改档位/审批通道）。
func subAgentTester(t *testing.T, prov *scriptedProvider, maxDepth, conc int, approveAll bool) (*Agent, *tools.Tool) {
	t.Helper()
	a, _ := newTestAgent(t)
	a.Provider = prov
	a.MaxDepth = maxDepth
	a.SubAgentConcurrency = conc
	a.Opt.ApproveAll = approveAll
	if err := a.Tools.Register(a.SubAgentTool()); err != nil {
		t.Fatal(err)
	}
	tl, ok := a.Tools.Get("agent")
	if !ok {
		t.Fatal("agent tool not registered")
	}
	return a, tl
}

// TestSubAgentSessionResume：结果里带 sessionID，带回去可以接着聊；
// 第二次调用必须看得到上一轮的提问与回答，而不是从空上下文重来。
func TestSubAgentSessionResume(t *testing.T) {
	calls := 0
	var secondMsgs []llm.Message
	prov := &scriptedProvider{fn: func(_ int, _ string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		calls++
		if calls == 1 {
			return llm.Message{Role: llm.RoleAssistant, Content: "FIRST"}, nil
		}
		secondMsgs = append([]llm.Message(nil), msgs...)
		return llm.Message{Role: llm.RoleAssistant, Content: "SECOND"}, nil
	}}
	_, tl := subAgentTester(t, prov, 1, 1, true)

	res, err := tl.Run(context.Background(), `{"prompt":"child task"}`)
	if err != nil {
		t.Fatal(err)
	}
	m := regexp.MustCompile(`sessionID="([^"]+)"`).FindStringSubmatch(res.Text)
	if m == nil {
		t.Fatalf("result has no sessionID: %q", res.Text)
	}
	id := m[1]

	res2, err := tl.Run(context.Background(), fmt.Sprintf(`{"prompt":"dig deeper","session_id":%q}`, id))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(res2.Text, "SECOND") || !strings.Contains(res2.Text, id) {
		t.Fatalf("resume result: %q", res2.Text)
	}
	var sawTask, sawFirst bool
	for _, msg := range secondMsgs {
		if strings.Contains(msg.Content, "child task") {
			sawTask = true
		}
		if strings.Contains(msg.Content, "FIRST") {
			sawFirst = true
		}
	}
	if !sawTask || !sawFirst {
		t.Fatalf("resumed sub-agent did not see its history: %+v", secondMsgs)
	}
}

// TestSubAgentUnknownSession：进程里没有那个子会话时要说清楚，而不是悄悄开新的。
func TestSubAgentUnknownSession(t *testing.T) {
	prov := &scriptedProvider{fn: func(int, string, []llm.Message, []llm.ToolDef) (llm.Message, error) {
		return llm.Message{Role: llm.RoleAssistant, Content: "x"}, nil
	}}
	_, tl := subAgentTester(t, prov, 1, 1, true)
	_, err := tl.Run(context.Background(), `{"prompt":"x","session_id":"sa-999"}`)
	if err == nil || !strings.Contains(err.Error(), "no such subagent session") {
		t.Fatalf("want unknown-session error, got %v", err)
	}
}

// TestSubAgentDepthLadder：深度逐层递增，且每层重新绑定自己的 agent 工具
// （否则子 agent 手里的 agent 仍指向顶层，MaxDepth 失效）。
func TestSubAgentDepthLadder(t *testing.T) {
	var childHasAgent, grandchildHasAgent bool
	prov := &scriptedProvider{fn: func(call int, _ string, _ []llm.Message, defs []llm.ToolDef) (llm.Message, error) {
		switch call {
		case 1: // 子 agent（depth 1 < MaxDepth 2）：应能再派发
			childHasAgent = hasName(defs, "agent")
			return llm.Message{Role: llm.RoleAssistant, ToolCalls: []llm.ToolCall{
				{ID: "a2", Name: "agent", Input: `{"prompt":"grandchild"}`},
			}}, nil
		case 2: // 孙 agent（depth 2 == MaxDepth 2）：不能再派发
			grandchildHasAgent = hasName(defs, "agent")
			return llm.Message{Role: llm.RoleAssistant, Content: "GRAND-FINAL"}, nil
		default: // 子 agent 拿到孙结果后的最终回复
			return llm.Message{Role: llm.RoleAssistant, Content: "CHILD-FINAL"}, nil
		}
	}}

	_, tl := subAgentTester(t, prov, 2, 2, true)
	res, err := tl.Run(context.Background(), `{"prompt":"child"}`)
	if err != nil {
		t.Fatal(err)
	}
	if !childHasAgent {
		t.Error("child (depth 1 < max 2) should be able to delegate further")
	}
	if grandchildHasAgent {
		t.Error("grandchild (depth 2 == max 2) must not have the agent tool")
	}
	if !strings.Contains(res.Text, "CHILD-FINAL") {
		t.Errorf("child final reply not returned: %q", res.Text)
	}
}

// TestSubAgentDepthZeroDisables：深度上限为 0 时工具直接拒绝（零值 fail-closed）。
func TestSubAgentDepthZeroDisables(t *testing.T) {
	prov := &scriptedProvider{fn: func(int, string, []llm.Message, []llm.ToolDef) (llm.Message, error) {
		t.Error("provider must not be called when depth is 0")
		return llm.Message{}, nil
	}}
	_, tl := subAgentTester(t, prov, 0, 2, true)
	if _, err := tl.Run(context.Background(), `{"prompt":"x"}`); err == nil {
		t.Fatal("depth 0 must refuse to spawn a sub-agent")
	}
}

// taskProvider 按任务名（消息末条内容）返回结果，并记录并发峰值。
func taskProvider(delay time.Duration, failOn string, peak *int64) *scriptedProvider {
	var inflight, maxSeen int64
	return &scriptedProvider{fn: func(_ int, _ string, msgs []llm.Message, _ []llm.ToolDef) (llm.Message, error) {
		cur := atomic.AddInt64(&inflight, 1)
		for {
			old := atomic.LoadInt64(&maxSeen)
			if cur <= old || atomic.CompareAndSwapInt64(&maxSeen, old, cur) {
				break
			}
		}
		defer atomic.AddInt64(&inflight, -1)

		var last string
		if len(msgs) > 0 {
			last = msgs[len(msgs)-1].Content
		}
		if failOn != "" && strings.Contains(last, failOn) {
			return llm.Message{}, fmt.Errorf("boom on %s", last)
		}
		time.Sleep(delay)
		if peak != nil {
			seen := atomic.LoadInt64(&maxSeen) // 必须原子读：其他 goroutine 在原子写
			for {
				old := atomic.LoadInt64(peak)
				if seen <= old || atomic.CompareAndSwapInt64(peak, old, seen) {
					break
				}
			}
		}
		return llm.Message{Role: llm.RoleAssistant, Content: "R-" + strings.TrimSpace(last)}, nil
	}}
}

const subDelay = 150 * time.Millisecond

// TestSubAgentTasksParallel：ApproveAll 下多任务并行（总耗时接近单任务）。
func TestSubAgentTasksParallel(t *testing.T) {
	_, tl := subAgentTester(t, taskProvider(subDelay, "", nil), 1, 2, true)

	start := time.Now()
	res, err := tl.Run(context.Background(), `{"tasks":["T1","T2"]}`)
	elapsed := time.Since(start)
	if err != nil {
		t.Fatal(err)
	}
	if elapsed >= 2*subDelay {
		t.Errorf("sub-tasks look serialized: elapsed=%v", elapsed)
	}
	if !strings.Contains(res.Text, "R-T1") || !strings.Contains(res.Text, "R-T2") {
		t.Errorf("both sub-results expected: %q", res.Text)
	}
	// 汇总按任务顺序编号。
	if strings.Index(res.Text, "[1/2]") > strings.Index(res.Text, "[2/2]") {
		t.Errorf("summary should be ordered: %q", res.Text)
	}
}

// TestSubAgentTasksSerialWhenApproverExists：只要问得到人（挂了审批回调，或真有
// stdin），多任务就串行——危险门不看 ApproveAll，而宿主只有一个待裁决槽位，并行会
// 让前一条询问挂到超时。真无人值守（两个通道都没有）才并行，见下一个用例。
func TestSubAgentTasksSerialWhenApproverExists(t *testing.T) {
	a, tl := subAgentTester(t, taskProvider(subDelay, "", nil), 1, 2, true)
	a.Opt.Approval = func(ApprovalRequest) ApprovalVerdict { return ApprovalOnce }

	start := time.Now()
	if _, err := tl.Run(context.Background(), `{"tasks":["T1","T2"]}`); err != nil {
		t.Fatal(err)
	}
	if elapsed := time.Since(start); elapsed < 2*subDelay {
		t.Errorf("有审批通道时必须串行: elapsed=%v", elapsed)
	}
}

// TestSubAgentTasksParallelWhenNoApprover：ApproveAll 且没有任何审批通道
// （-p / cron / 渠道）才是并行的正当场景，即使档位不是 bypass。
func TestSubAgentTasksParallelWhenNoApprover(t *testing.T) {
	_, tl := subAgentTester(t, taskProvider(subDelay, "", nil), 1, 2, false)

	start := time.Now()
	if _, err := tl.Run(context.Background(), `{"tasks":["T1","T2"]}`); err != nil {
		t.Fatal(err)
	}
	if elapsed := time.Since(start); elapsed >= 2*subDelay {
		t.Errorf("无人可问时应并行: elapsed=%v", elapsed)
	}
}

// TestSubAgentTasksFailureIsolation：单个子任务失败不拖垮整批，工具整体仍算成功。
func TestSubAgentTasksFailureIsolation(t *testing.T) {
	_, tl := subAgentTester(t, taskProvider(0, "BAD", nil), 1, 2, true)

	res, err := tl.Run(context.Background(), `{"tasks":["GOOD","BAD"]}`)
	if err != nil {
		t.Fatalf("batch with one failure should not fail the whole tool: %v", err)
	}
	if !strings.Contains(res.Text, "R-GOOD") {
		t.Errorf("successful task result missing: %q", res.Text)
	}
	if !strings.Contains(res.Text, "[2/2] failed") {
		t.Errorf("failed task should be marked: %q", res.Text)
	}
}

// TestSubAgentAllTasksFail：全部失败时以 Result.Err 返回（同 bash 约定：保留文本
// 同时表明失败），让父级失败熔断计数生效。
func TestSubAgentAllTasksFail(t *testing.T) {
	_, tl := subAgentTester(t, taskProvider(0, "T", nil), 1, 2, true)
	res, err := tl.Run(context.Background(), `{"tasks":["T1","T2"]}`)
	if err != nil {
		t.Fatalf("partial output should ride in Result, not the error slot: %v", err)
	}
	if res.Err == nil {
		t.Fatal("all-failed batch must report failure via Result.Err")
	}
	// 文本仍然保留各任务失败原因，便于父模型判断。
	if !strings.Contains(res.Text, "failed") {
		t.Errorf("failure detail should be preserved: %q", res.Text)
	}
}

// TestSubAgentConcurrencyCap：并发度不超过 SubAgentConcurrency。
func TestSubAgentConcurrencyCap(t *testing.T) {
	var peak int64
	_, tl := subAgentTester(t, taskProvider(40*time.Millisecond, "", &peak), 1, 2, true)

	if _, err := tl.Run(context.Background(), `{"tasks":["A","B","C","D","E"]}`); err != nil {
		t.Fatal(err)
	}
	if peak > 2 {
		t.Errorf("in-flight sub-agents = %d, want <= 2", peak)
	}
	if peak < 2 {
		t.Errorf("expected the batch to actually run in parallel, peak = %d", peak)
	}
}

// TestSubAgentInputValidation：prompt/tasks 互斥、空项与超量都要拒绝。
func TestSubAgentInputValidation(t *testing.T) {
	_, tl := subAgentTester(t, taskProvider(0, "", nil), 1, 2, true)
	ctx := context.Background()
	for _, in := range []string{
		`{}`,
		`{"prompt":"a","tasks":["b"]}`,
		`{"tasks":[]}`,
		`{"tasks":["ok","  "]}`,
		`{"tasks":["1","2","3","4","5","6","7","8","9","10","11","12","13","14","15","16","17"]}`,
	} {
		if _, err := tl.Run(ctx, in); err == nil {
			t.Errorf("input %s should be rejected", in)
		}
	}
}