package run

import (
	"encoding/json"
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/workflow/invocation"
	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
	"github.com/tidbcloud/AgentTeam/internal/workflow/runlimit"
)

func TestFanOutInstructionReferenceFailureIsCollectedAndBlocksSuccessors(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, numericFanOutRuntimePlan())
	root := workItemRaw("root", nil, "@dev1")
	child := workItemRaw("child", []string{"root"}, "@dev2")
	child[plan.WorkItemFieldInstruction] = "use $root.missing"
	grandchild := workItemRaw("grandchild", []string{"child"}, "@dev1")
	state := startFanOut(t, kernel, []any{root, child, grandchild})

	rootIntent, state := stepToMessageIntent(t, kernel, state)
	transition, err := kernel.Complete(state, Completion{
		ActivationID: rootIntent.ActivationID,
		Port:         "completed",
		Output: Values{
			"answer":    json.Number("42"),
			"workScope": messageScopeObject(*rootIntent.WorkScope),
		},
	}, reducerNow)
	require.NoError(t, err, "a bad dynamic reference is a collected WorkItem failure, not a stuck reducer error")
	scope := onlyFanOutScope(t, transition.State)
	failed := fanOutItemByKey(scope.FanOut, "child")
	require.Equal(t, WorkItemFailed, failed.Status)
	require.Equal(t, workItemInstructionReferenceFailureCode, failed.Failure.Code)
	require.Contains(t, failed.Failure.Message, "does not exist")
	require.Equal(t, "work_item", failed.Failure.Stage)
	require.NotNil(t, failed.Failure.WorkItem)
	require.Equal(t, "child", failed.Failure.WorkItem.Key)
	require.NotNil(t, failed.Failure.WorkItem.Index)
	require.Equal(t, failed.Index, *failed.Failure.WorkItem.Index)
	require.Empty(t, failed.RootActivationID, "no effect-producing child activation is created")
	blocked := fanOutItemByKey(scope.FanOut, "grandchild")
	require.Equal(t, WorkItemBlocked, blocked.Status)
	require.Equal(t, []string{"child"}, blocked.BlockedBy)
	require.Equal(t, ScopeSucceeded, scope.Status, "failureMode=collect resolves after the blocked propagation")
}

func TestFanOutNonNormalizedNumberBudgetFailureIsCollected(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, numericFanOutRuntimePlan())
	root := workItemRaw("root", nil, "@dev1")
	child := workItemRaw("child", []string{"root"}, "@dev2")
	child[plan.WorkItemFieldInstruction] = "$root.answer"
	state := startFanOut(t, kernel, []any{root, child})

	rootIntent, state := stepToMessageIntent(t, kernel, state)
	transition, err := kernel.Complete(state, Completion{
		ActivationID: rootIntent.ActivationID,
		Port:         "completed",
		Output: Values{
			"answer":    json.Number("1e4095"),
			"workScope": messageScopeObject(*rootIntent.WorkScope),
		},
	}, reducerNow)
	require.NoError(t, err)
	scope := onlyFanOutScope(t, transition.State)
	failed := fanOutItemByKey(scope.FanOut, "root")
	require.Equal(t, WorkItemFailed, failed.Status)
	require.Equal(t, workItemResourceLimitFailureCode, failed.Failure.Code)
	require.Equal(t, WorkItemBlocked, fanOutItemByKey(scope.FanOut, "child").Status)
	activation, activationErr := transition.State.activation(rootIntent.ActivationID)
	require.NoError(t, activationErr)
	require.Nil(t, activation.Output)
	require.Equal(t, ScopeSucceeded, scope.Status)
}

func TestFanOutFixedTaskCallbackFreezesResponderAndProjectsCausalWatchHandoff(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	sharedTask := map[string]any{"kind": "task", "scopeId": "tsk_shared"}
	items := []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child", []string{"root"}, "@dev2"),
		workItemRaw("grandchild", []string{"child"}, "@dev1"),
		workItemRaw("independent", nil, "@dev2"),
	}
	for _, raw := range items {
		raw.(map[string]any)[plan.WorkItemFieldWorkScope] = sharedTask
	}
	state := startFanOut(t, kernel, items)
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, rootIntent, "agt_root")
	independentIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, independentIntent, "agt_child")
	require.NotNil(t, rootIntent.TaskWatchHandoff)
	require.Equal(t, "tsk_shared", rootIntent.TaskWatchHandoff.TaskScope.ScopeID)
	rootItem := fanOutItemByKey(onlyFanOutScope(t, state).FanOut, "root")
	require.Equal(
		t,
		rootItem.ID,
		rootIntent.TaskWatchHandoff.WorkItemID,
		"a fan-out callback retains its concrete WorkItem identity rather than the static-callback fallback",
	)
	require.Empty(t, rootIntent.TaskWatchHandoff.CausalPredecessorAgentIDs)
	require.NotNil(t, independentIntent.TaskWatchHandoff)

	state = completeIntentAsResponder(t, kernel, state, rootIntent, "agt_root", Values{"answer": "root"})
	rootItem = fanOutItemByKey(onlyFanOutScope(t, state).FanOut, "root")
	require.Equal(t, []string{"agt_root"}, rootItem.ActivatedResponderAgentIDs)

	childIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, []string{"agt_root"}, childIntent.TaskWatchHandoff.CausalPredecessorAgentIDs)
	state = completeIntentAsResponder(t, kernel, state, childIntent, "agt_child", Values{"answer": "child"})

	grandchildIntent, state := stepToMessageIntent(t, kernel, state)
	require.Empty(
		t,
		grandchildIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"agt_root was already replaced by child and the independent @dev2 branch keeps agt_child on the frontier",
	)
	require.NotContains(
		t,
		grandchildIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"agt_independent",
		"an activated but non-causal WorkItem must never be demoted",
	)

	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	require.Equal(t, state, restored, "frozen resolved responder IDs survive recovery")
}

func TestFanOutTaskWatchHandoffOnlyDemotesTheCurrentCausalFrontier(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	sharedTask := map[string]any{"kind": "task", "scopeId": "tsk_shared"}
	items := []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child", []string{"root"}, "@dev2"),
		workItemRaw("grandchild", []string{"child"}, "@dev1"),
	}
	for _, raw := range items {
		raw.(map[string]any)[plan.WorkItemFieldWorkScope] = sharedTask
	}

	state := startFanOut(t, kernel, items)
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntentAsResponder(t, kernel, state, rootIntent, "agt_root", Values{"answer": "root"})
	childIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, []string{"agt_root"}, childIntent.TaskWatchHandoff.CausalPredecessorAgentIDs)

	state = completeIntentAsResponder(t, kernel, state, childIntent, "agt_child", Values{"answer": "child"})
	grandchildIntent, _ := stepToMessageIntent(t, kernel, state)
	require.Equal(
		t,
		[]string{"agt_child"},
		grandchildIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"the already-superseded root is not part of the current routing frontier",
	)
}

func TestFanOutTaskWatchHandoffExcludesCausalPredecessorFromDifferentWorkScope(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	state := startFanOut(t, kernel, []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child", []string{"root"}, "@dev2"),
	})
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntentAsResponder(t, kernel, state, rootIntent, "agt_root", Values{"answer": "root"})
	childIntent, _ := stepToMessageIntent(t, kernel, state)
	require.NotNil(t, childIntent.TaskWatchHandoff)
	require.Empty(
		t,
		childIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"causal WorkItems in a different normalized Task must not be demoted",
	)
}

func TestFanOutTaskWatchHandoffKeepsSharedAgentUntilEveryParallelBranchIsReplaced(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	sharedTask := map[string]any{"kind": "task", "scopeId": "tsk_shared"}
	items := []any{
		workItemRaw("root_a", nil, "@dev1"),
		workItemRaw("root_b", nil, "@dev2"),
		workItemRaw("child_a", []string{"root_a"}, "@dev2"),
		workItemRaw("review", []string{"child_a", "root_b"}, "@dev1"),
	}
	for _, raw := range items {
		raw.(map[string]any)[plan.WorkItemFieldWorkScope] = sharedTask
	}

	state := startFanOut(t, kernel, items)
	rootAIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, rootAIntent, "agt_shared")
	rootBIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, rootBIntent, "agt_shared")
	require.Equal(t, "root_a", activationItemKey(t, state, rootAIntent.ActivationID))
	require.Equal(t, "root_b", activationItemKey(t, state, rootBIntent.ActivationID))

	state = completeIntentAsResponder(t, kernel, state, rootAIntent, "agt_shared", Values{"answer": "a"})
	state = completeIntentAsResponder(t, kernel, state, rootBIntent, "agt_shared", Values{"answer": "b"})
	childIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, "child_a", activationItemKey(t, state, childIntent.ActivationID))
	require.Empty(
		t,
		childIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"the normalized Agent remains on root_b even though the two branches used different responder refs",
	)

	state = completeIntentAsResponder(t, kernel, state, childIntent, "agt_child", Values{"answer": "child"})
	reviewIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, "review", activationItemKey(t, state, reviewIntent.ActivationID))
	require.Equal(
		t,
		[]string{"agt_shared", "agt_child"},
		reviewIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"a join successor replaces both causal branch frontiers and deduplicates their shared Agent",
	)
}

func TestFanOutTaskWatchHandoffKeepsForkPredecessorUntilEveryBranchActivates(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	sharedTask := map[string]any{"kind": "task", "scopeId": "tsk_shared"}
	items := []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child_a", []string{"root"}, "@dev2"),
		workItemRaw("child_b", []string{"root"}, "@dev2"),
		workItemRaw("review", []string{"child_a", "child_b"}, "@dev1"),
	}
	for _, raw := range items {
		raw.(map[string]any)[plan.WorkItemFieldWorkScope] = sharedTask
	}

	state := startFanOut(t, kernel, items)
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntentAsResponder(t, kernel, state, rootIntent, "agt_root", Values{"answer": "root"})

	childAIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, childAIntent, "agt_child")
	require.Equal(t, "child_a", activationItemKey(t, state, childAIntent.ActivationID))
	require.Empty(
		t,
		childAIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"root still represents the not-yet-activated child_b branch",
	)

	childBIntent, state := stepToMessageIntent(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, childBIntent, "agt_child")
	require.Equal(t, "child_b", activationItemKey(t, state, childBIntent.ActivationID))
	require.Equal(
		t,
		[]string{"agt_root"},
		childBIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"the second successor covers root's final outstanding branch",
	)

	state = completeIntentAsResponder(t, kernel, state, childAIntent, "agt_child", Values{"answer": "a"})
	state = completeIntentAsResponder(t, kernel, state, childBIntent, "agt_child", Values{"answer": "b"})
	reviewIntent, _ := stepToMessageIntent(t, kernel, state)
	require.Equal(
		t,
		[]string{"agt_child"},
		reviewIntent.TaskWatchHandoff.CausalPredecessorAgentIDs,
		"the join successor replaces both child branches and deduplicates their shared Agent",
	)
}

func TestFanOutPropagatesRootBlockersTransitivelyAndKeepsIndependentWorkRunning(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	state := startFanOut(t, kernel, []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child", []string{"root"}, "@dev2"),
		workItemRaw("grandchild", []string{"child"}, "@dev1"),
		workItemRaw("independent", nil, "@dev2"),
	})
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	independentIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntent(t, kernel, state, rootIntent, "failed", Values{}, &Failure{Code: "root.failed"})
	scope := onlyFanOutScope(t, state)
	require.Equal(
		t,
		[]WorkItemStatus{WorkItemFailed, WorkItemBlocked, WorkItemBlocked, WorkItemRunning},
		itemStatuses(scope),
	)
	require.Equal(t, []string{"root"}, scope.FanOut.Items[1].BlockedBy)
	require.Equal(t, []string{"root"}, scope.FanOut.Items[2].BlockedBy)

	state = completeIntent(t, kernel, state, independentIntent, "completed", Values{"answer": "still ran"}, nil)
	scope = onlyFanOutScope(t, state)
	require.Equal(t, ScopeSucceeded, scope.Status)
	terminal := stepUntilTerminal(t, kernel, state, 4)
	require.Equal(t, StatusSucceeded, terminal.Status)
}

func TestFanOutMaximumDependencyCascadeWithLongKeysRemainsPersistable(t *testing.T) {
	t.Parallel()
	const (
		rootCount      = 128
		dependentCount = plan.MaxWorkItemDependencyEdges / rootCount
	)
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	longKey := func(kind string, index int) string {
		prefix := fmt.Sprintf("%s_%03d_", kind, index)
		return prefix + strings.Repeat("k", plan.MaxWorkItemKeyBytes-len(prefix))
	}
	item := func(key string, dependencies []string, scopeIndex int) map[string]any {
		copiedDependencies := make([]any, len(dependencies))
		for index, dependency := range dependencies {
			copiedDependencies[index] = dependency
		}
		return map[string]any{
			plan.WorkItemFieldKey:           key,
			plan.WorkItemFieldDependsOn:     copiedDependencies,
			plan.WorkItemFieldWorkScope:     map[string]any{"kind": "task", "scopeId": fmt.Sprintf("tsk_dense_%03d", scopeIndex)},
			plan.WorkItemFieldRelatedScopes: []any{map[string]any{"kind": "task", "scopeId": "tsk_parent"}},
			plan.WorkItemFieldResponder:     "@dev1",
			plan.WorkItemFieldInstruction:   "check dense dependency settlement",
		}
	}

	rootKeys := make([]string, rootCount)
	workItems := make([]any, 0, rootCount+dependentCount)
	for index := range rootKeys {
		rootKeys[index] = longKey("root", index)
		require.Len(t, []byte(rootKeys[index]), plan.MaxWorkItemKeyBytes)
		workItems = append(workItems, item(rootKeys[index], nil, index))
	}
	for index := 0; index < dependentCount; index++ {
		workItems = append(workItems, item(longKey("dependent", index), rootKeys, rootCount+index))
	}

	state := startFanOut(t, kernel, workItems)
	scope := onlyFanOutScope(t, state)
	edges := 0
	for index := range scope.FanOut.Items {
		edges += len(scope.FanOut.Items[index].DependsOn)
		if index >= rootCount {
			continue
		}
		itemState := &scope.FanOut.Items[index]
		projectCompactFanOutResourceFailure(itemState)
		if itemState.RootActivationID == "" {
			continue
		}
		activation, err := state.activation(itemState.RootActivationID)
		require.NoError(t, err)
		activation.Status = ActivationFailed
		activation.Failure = &Failure{Code: workItemResourceLimitFailureCode, Message: workItemResourceLimitMessage}
		activation.Output = nil
		activation.SelectedPorts = []string{}
	}
	require.Equal(t, plan.MaxWorkItemDependencyEdges, edges)

	transition := &Transition{State: state}
	changed, err := kernel.blockUnrunnableItems(transition, scope)
	require.NoError(t, err, "the maximum legal dependency cascade must settle item-locally")
	require.True(t, changed)
	for index := rootCount; index < len(scope.FanOut.Items); index++ {
		blocked := &scope.FanOut.Items[index]
		require.Equal(t, WorkItemBlocked, blocked.Status)
		require.Len(t, blocked.BlockedBy, rootCount)
		require.Equal(t, rootKeys, blocked.BlockedBy)
	}
	require.NoError(t, validateFanOutHardStateBudget("dense long-key dependency cascade", state))
	encoded, err := json.Marshal(state)
	require.NoError(t, err)
	_, err = invocation.DeterministicJSONObjectWithBudget(
		"workflow Run state", encoded,
		runlimit.MaxRunStateBytes,
		runlimit.MaxRunStateJSONNesting,
		runlimit.MaxRunStateJSONValues,
	)
	require.NoError(t, err)
}

func TestFanOutRecoveryCancelAndLateCompletionAreFenced(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	state := startFanOut(t, kernel, []any{
		workItemRaw("a", nil, "@dev1"),
		workItemRaw("b", nil, "@dev2"),
		workItemRaw("c", nil, "@dev1"),
	})
	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	require.Equal(t, state, restored)

	first, err := kernel.Step(restored, reducerNow)
	require.NoError(t, err)
	replayedFromSameSnapshot, err := kernel.Step(restored, reducerNow)
	require.NoError(t, err)
	require.Equal(t, first, replayedFromSameSnapshot, "recovery replay is byte-semantic deterministic")
	require.Len(t, first.Intents, 1)

	canceled, err := kernel.Cancel(first.State)
	require.NoError(t, err)
	scope := onlyFanOutScope(t, canceled.State)
	require.Equal(t, ScopeCanceled, scope.Status)
	require.Equal(
		t,
		[]WorkItemStatus{WorkItemCanceled, WorkItemCanceled, WorkItemCanceled},
		itemStatuses(scope),
	)
	replayedCancel, err := kernel.Cancel(canceled.State)
	require.NoError(t, err)
	require.Equal(t, canceled.State, replayedCancel.State)

	_, err = kernel.Complete(canceled.State, Completion{
		ActivationID: first.Intents[0].ActivationID,
		Port:         "completed",
		Output: Values{
			"answer":    "late",
			"workScope": messageScopeObject(*first.Intents[0].WorkScope),
		},
	}, reducerNow)
	require.ErrorContains(t, err, "terminal")

	corrupt := state.Clone()
	corrupt.Scopes[0].FanOut.Items[0].ID = "forged"
	_, err = kernel.Step(corrupt, reducerNow)
	require.ErrorContains(t, err, "invalid semantic ID")
}