package run

import (
	"strings"
	"testing"

	"github.com/stretchr/testify/require"
)

func TestTaskWatchFrontierSerialSameAndDifferentTasks(t *testing.T) {
	t.Parallel()
	for _, test := range []struct {
		name             string
		secondTaskScope  string
		wantPredecessors []string
	}{
		{
			name:             "same Task hands off to the successor",
			secondTaskScope:  "tsk_shared",
			wantPredecessors: []string{"agt_a"},
		},
		{
			name:             "different Tasks keep independent frontiers",
			secondTaskScope:  "tsk_other",
			wantPredecessors: []string{},
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			kernel := compilePlanKernel(t, taskWatchSerialRuntimePlan)
			state, err := kernel.Start(Values{
				"task_a":  taskWatchScopeObject("tsk_shared"),
				"task_b":  taskWatchScopeObject(test.secondTaskScope),
				"agent_a": "@agent-a",
				"agent_b": "@agent-b",
			})
			require.NoError(t, err)

			first, state := stepToNextTaskWatchCallback(t, kernel, state)
			require.Equal(t, "a", first.NodeID)
			requireTaskWatchHandoff(t, first, "tsk_shared", []string{})
			require.Equal(t, first.ActivationID, first.TaskWatchHandoff.WorkItemID)
			state = completeTaskWatchCallback(
				t, kernel, state, first, "completed", "agt_a", Values{"value": "a"},
			)

			second, _ := stepToNextTaskWatchCallback(t, kernel, state)
			require.Equal(t, "b", second.NodeID)
			requireTaskWatchHandoff(t, second, test.secondTaskScope, test.wantPredecessors)
			require.Equal(t, second.ActivationID, second.TaskWatchHandoff.WorkItemID)
		})
	}
}

func TestTaskWatchFrontierIndependentBranchesAndJoinSuccessor(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchIndependentBranchesRuntimePlan)
	state, err := kernel.Start(taskWatchParallelInput())
	require.NoError(t, err)

	a, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, a, "agt_a")
	c, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "a", a.NodeID)
	require.Equal(t, "c", c.NodeID)
	requireTaskWatchHandoff(t, a, "tsk_shared", []string{})
	requireTaskWatchHandoff(t, c, "tsk_shared", []string{})

	state = completeTaskWatchCallback(t, kernel, state, a, "completed", "agt_a", Values{"value": "a"})
	state = completeTaskWatchCallback(t, kernel, state, c, "completed", "agt_c", Values{"value": "c"})

	b, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, b, "agt_b")
	d, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "b", b.NodeID)
	require.Equal(t, "d", d.NodeID)
	requireTaskWatchHandoff(t, b, "tsk_shared", []string{"agt_a"})
	requireTaskWatchHandoff(t, d, "tsk_shared", []string{"agt_c"})

	state = completeTaskWatchCallback(t, kernel, state, b, "completed", "agt_b", Values{"value": "b"})
	state = completeTaskWatchCallback(t, kernel, state, d, "completed", "agt_d", Values{"value": "d"})

	e, _ := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "e", e.NodeID)
	requireTaskWatchHandoff(t, e, "tsk_shared", []string{"agt_b", "agt_d"})
}

func TestTaskWatchFrontierCommonPredecessorForkAndJoinSuccessor(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchCommonPredecessorRuntimePlan)
	state, err := kernel.Start(taskWatchParallelInput())
	require.NoError(t, err)

	a, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "a", a.NodeID)
	requireTaskWatchHandoff(t, a, "tsk_shared", []string{})
	state = completeTaskWatchCallback(t, kernel, state, a, "completed", "agt_a", Values{"value": "a"})

	b, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, b, "agt_b")
	c, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "b", b.NodeID)
	require.Equal(t, "c", c.NodeID)
	requireTaskWatchHandoff(
		t,
		b,
		"tsk_shared",
		[]string{},
	)
	requireTaskWatchHandoff(
		t,
		c,
		"tsk_shared",
		[]string{"agt_a"},
	)

	state = completeTaskWatchCallback(t, kernel, state, b, "completed", "agt_b", Values{"value": "b"})
	state = completeTaskWatchCallback(t, kernel, state, c, "completed", "agt_c", Values{"value": "c"})

	e, _ := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "e", e.NodeID)
	requireTaskWatchHandoff(t, e, "tsk_shared", []string{"agt_b", "agt_c"})
}

func TestRestoreUndeliveredCallbackTaskWatchRestoresOnlyProjection(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchCommonPredecessorRuntimePlan)
	state, err := kernel.Start(taskWatchParallelInput())
	require.NoError(t, err)

	predecessor, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, predecessor, "agt_a")
	state = completeTaskWatchCallback(
		t, kernel, state, predecessor, "completed", "agt_a", Values{"value": "a"},
	)

	projected, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "b", projected.NodeID)
	projectedActivation, err := state.activation(projected.ActivationID)
	require.NoError(t, err)
	require.Equal(t, []string{projected.ActivationID}, projectedActivation.TaskWatchFrontier)
	require.Equal(
		t,
		[]string{predecessor.ActivationID},
		projectedActivation.CallbackTaskWatch.ReplacedActivationIDs,
	)

	restored, err := kernel.RestoreUndeliveredCallbackTaskWatch(state, projected.ActivationID)
	require.NoError(t, err)
	require.Equal(t, []string{"callback_task_watch_projection_restored"}, eventKinds(restored.Events))
	restoredActivation, err := restored.State.activation(projected.ActivationID)
	require.NoError(t, err)
	require.Equal(t, ActivationWaiting, restoredActivation.Status,
		"bookkeeping restoration must not choose the callback failure port")
	require.Equal(t, []string{predecessor.ActivationID}, restoredActivation.TaskWatchFrontier)
	require.Empty(t, restoredActivation.CallbackTaskWatch.ResponderAgentID)

	replay, err := kernel.RestoreUndeliveredCallbackTaskWatch(restored.State, projected.ActivationID)
	require.NoError(t, err)
	require.Empty(t, replay.Events)
	require.Equal(t, restored.State, replay.State)
}

func TestRestoreUndeliveredCallbackTaskWatchDoesNotResurrectMovedProjection(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchCommonPredecessorRuntimePlan)
	state, err := kernel.Start(taskWatchParallelInput())
	require.NoError(t, err)

	predecessor, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, predecessor, "agt_a")
	state = completeTaskWatchCallback(
		t, kernel, state, predecessor, "completed", "agt_a", Values{"value": "a"},
	)

	projected, state := stepToNextTaskWatchCallback(t, kernel, state)
	projectedActivation, err := state.activation(projected.ActivationID)
	require.NoError(t, err)
	require.Equal(t, []string{projected.ActivationID}, projectedActivation.TaskWatchFrontier)

	// Model a future caller that reaches the old callback only after its
	// provisional token has already moved to another control carrier. Restoring
	// here must be a no-op; otherwise the predecessor would be resurrected as a
	// second live branch beside the moved token.
	var downstream *Activation
	for index := range state.Activations {
		if state.Activations[index].ID != projected.ActivationID && state.Activations[index].NodeID == "c" {
			downstream = &state.Activations[index]
			break
		}
	}
	require.NotNil(t, downstream)
	projectedActivation.TaskWatchFrontier = nil
	downstream.TaskWatchFrontier = []string{projected.ActivationID}
	before := state.Clone()

	restored, err := kernel.RestoreUndeliveredCallbackTaskWatch(state, projected.ActivationID)
	require.NoError(t, err)
	require.Empty(t, restored.Events)
	require.Equal(t, before, restored.State)
}

func TestTaskWatchFrontierProtectsResolvedAgentOnAnotherBranch(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchIndependentBranchesRuntimePlan)
	state, err := kernel.Start(taskWatchParallelInput())
	require.NoError(t, err)

	a, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, a, "agt_shared")
	c, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = completeTaskWatchCallback(t, kernel, state, a, "completed", "agt_shared", Values{"value": "a"})
	state = completeTaskWatchCallback(t, kernel, state, c, "completed", "agt_shared", Values{"value": "c"})

	b, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = observeTaskWatchResponder(t, kernel, state, b, "agt_b")
	require.Equal(t, "b", b.NodeID)
	requireTaskWatchHandoff(
		t,
		b,
		"tsk_shared",
		[]string{},
	)

	d, _ := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "d", d.NodeID)
	requireTaskWatchHandoff(t, d, "tsk_shared", []string{"agt_shared"})
}

func TestTaskWatchFrontierCrossesFanOutChildScopesAndReturnsToParentTask(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchStaticFanOutRuntimePlan)
	state, err := kernel.Start(Values{
		"parent_task": taskWatchScopeObject("tsk_parent"),
		"planner":     "@planner",
		"calc_a":      "@calc-a",
		"calc_b":      "@calc-b",
		"reviewer":    "@reviewer",
	})
	require.NoError(t, err)

	plannerIntent, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "planner", plannerIntent.NodeID)
	requireTaskWatchHandoff(t, plannerIntent, "tsk_parent", []string{})
	state = completeTaskWatchCallback(
		t,
		kernel,
		state,
		plannerIntent,
		"completed",
		"agt_planner",
		Values{"workItems": []any{
			taskWatchWorkItem("calc_a", "tsk_child_a", "@calc-a"),
			taskWatchWorkItem("calc_b", "tsk_child_b", "@calc-b"),
		}},
	)

	calcAIntent, state := stepToNextTaskWatchCallback(t, kernel, state)
	calcBIntent, state := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "implement", calcAIntent.NodeID)
	require.Equal(t, "implement", calcBIntent.NodeID)
	requireTaskWatchHandoff(t, calcAIntent, "tsk_child_a", []string{})
	requireTaskWatchHandoff(t, calcBIntent, "tsk_child_b", []string{})
	require.True(t, strings.HasPrefix(calcAIntent.TaskWatchHandoff.WorkItemID, "wfi_"))
	require.True(t, strings.HasPrefix(calcBIntent.TaskWatchHandoff.WorkItemID, "wfi_"))
	require.NotEqual(t, calcAIntent.ActivationID, calcAIntent.TaskWatchHandoff.WorkItemID)
	require.NotEqual(t, calcBIntent.ActivationID, calcBIntent.TaskWatchHandoff.WorkItemID)

	state = completeTaskWatchCallback(
		t, kernel, state, calcAIntent, "completed", "agt_calc_a", Values{"answer": "11"},
	)
	state = completeTaskWatchCallback(
		t, kernel, state, calcBIntent, "completed", "agt_calc_b", Values{"answer": "22"},
	)

	reviewerIntent, _ := stepToNextTaskWatchCallback(t, kernel, state)
	require.Equal(t, "reviewer", reviewerIntent.NodeID)
	requireTaskWatchHandoff(t, reviewerIntent, "tsk_parent", []string{"agt_planner"})
	require.NotContains(t, reviewerIntent.TaskWatchHandoff.CausalPredecessorAgentIDs, "agt_calc_a")
	require.NotContains(t, reviewerIntent.TaskWatchHandoff.CausalPredecessorAgentIDs, "agt_calc_b")
}

func TestTaskWatchFrontierJSONReplayAndConflictingResponder(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchSerialRuntimePlan)
	state, err := kernel.Start(Values{
		"task_a":  taskWatchScopeObject("tsk_shared"),
		"task_b":  taskWatchScopeObject("tsk_shared"),
		"agent_a": "@agent-a",
		"agent_b": "@agent-b",
	})
	require.NoError(t, err)

	first, state := stepToNextTaskWatchCallback(t, kernel, state)
	state = completeTaskWatchCallback(
		t, kernel, state, first, "completed", "agt_a", Values{"value": "a"},
	)
	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	restoredAgain, err := JSONRoundTrip(restored)
	require.NoError(t, err)
	require.Equal(t, restored, restoredAgain, "Task Watch provenance must be stable after recovery normalization")
	originalFirst, err := state.activation(first.ActivationID)
	require.NoError(t, err)
	restoredFirst, err := restored.activation(first.ActivationID)
	require.NoError(t, err)
	require.Equal(t, originalFirst.CallbackTaskWatch, restoredFirst.CallbackTaskWatch)
	require.Equal(t, []string{first.ActivationID}, restored.Activations[1].TaskWatchFrontier)

	replayOutput := Values{
		"value":     "a",
		"workScope": taskWatchScopeObject("tsk_shared"),
	}
	replay, err := kernel.Complete(restored, Completion{
		ActivationID:             first.ActivationID,
		Port:                     "completed",
		Output:                   replayOutput,
		CallbackResponderAgentID: "agt_a",
	}, reducerNow)
	require.NoError(t, err)
	require.Empty(t, replay.Intents)
	require.Equal(t, []Event{{
		Kind: "completion_replayed", ActivationID: first.ActivationID, Port: "completed",
	}}, replay.Events)
	require.Equal(t, restored.Status, replay.State.Status)
	require.Equal(t, restored.NextActivation, replay.State.NextActivation)
	require.Len(t, replay.State.Activations, len(restored.Activations))
	replayedFirst, err := replay.State.activation(first.ActivationID)
	require.NoError(t, err)
	require.Equal(t, restoredFirst.CallbackTaskWatch, replayedFirst.CallbackTaskWatch)
	require.Equal(t, []string{first.ActivationID}, replay.State.Activations[1].TaskWatchFrontier)

	beforeConflictCallback := *restoredFirst.CallbackTaskWatch
	beforeConflictCallback.ReplacedActivationIDs = append(
		[]string{}, restoredFirst.CallbackTaskWatch.ReplacedActivationIDs...,
	)
	beforeConflictFrontier := append([]string{}, restored.Activations[1].TaskWatchFrontier...)
	beforeConflictActivationCount := len(restored.Activations)
	_, err = kernel.Complete(restored, Completion{
		ActivationID:             first.ActivationID,
		Port:                     "completed",
		Output:                   replayOutput,
		CallbackResponderAgentID: "agt_other",
	}, reducerNow)
	require.ErrorContains(t, err, "different result")
	require.Len(t, restored.Activations, beforeConflictActivationCount)
	afterConflictFirst, err := restored.activation(first.ActivationID)
	require.NoError(t, err)
	require.Equal(t, &beforeConflictCallback, afterConflictFirst.CallbackTaskWatch)
	require.Equal(t, beforeConflictFrontier, restored.Activations[1].TaskWatchFrontier)

	firstActivation, err := replay.State.activation(first.ActivationID)
	require.NoError(t, err)
	require.NotNil(t, firstActivation.CallbackTaskWatch)
	require.Equal(t, "agt_a", firstActivation.CallbackTaskWatch.ResponderAgentID)

	second, _ := stepToNextTaskWatchCallback(t, kernel, replay.State)
	requireTaskWatchHandoff(t, second, "tsk_shared", []string{"agt_a"})
}

func TestObserveCallbackResponderFreezesBeforeCompletionAndIsIdempotent(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, taskWatchSerialRuntimePlan)
	state, err := kernel.Start(Values{
		"task_a":  taskWatchScopeObject("tsk_shared"),
		"task_b":  taskWatchScopeObject("tsk_shared"),
		"agent_a": "@agent-a",
		"agent_b": "@agent-b",
	})
	require.NoError(t, err)

	intent, state := stepToNextTaskWatchCallback(t, kernel, state)
	before, err := state.activation(intent.ActivationID)
	require.NoError(t, err)
	require.Equal(t, ActivationWaiting, before.Status)
	require.Empty(t, before.CallbackTaskWatch.ResponderAgentID)

	observed, err := kernel.ObserveCallbackResponder(state, intent.ActivationID, "agt_a")
	require.NoError(t, err)
	require.Empty(t, observed.Intents)
	require.Equal(t, []Event{{
		Kind:         "callback_responder_observed",
		ActivationID: intent.ActivationID,
		NodeID:       intent.NodeID,
		AgentID:      "agt_a",
	}}, observed.Events)
	observedActivation, err := observed.State.activation(intent.ActivationID)
	require.NoError(t, err)
	require.Equal(t, ActivationWaiting, observedActivation.Status)
	require.Equal(t, "agt_a", observedActivation.CallbackTaskWatch.ResponderAgentID)

	replayed, err := kernel.ObserveCallbackResponder(observed.State, intent.ActivationID, "agt_a")
	require.NoError(t, err)
	require.Empty(t, replayed.Intents)
	require.Empty(t, replayed.Events)
	require.Equal(t, observed.State, replayed.State)

	_, err = kernel.ObserveCallbackResponder(replayed.State, intent.ActivationID, "agt_other")
	require.ErrorContains(t, err, "already frozen as")
	stillObserved, err := replayed.State.activation(intent.ActivationID)
	require.NoError(t, err)
	require.Equal(t, "agt_a", stillObserved.CallbackTaskWatch.ResponderAgentID)

	completed := completeTaskWatchCallback(
		t, kernel, replayed.State, intent, "completed", "agt_a", Values{"value": "a"},
	)
	completedActivation, err := completed.activation(intent.ActivationID)
	require.NoError(t, err)
	require.Equal(t, ActivationSucceeded, completedActivation.Status)
	require.Equal(t, "agt_a", completedActivation.CallbackTaskWatch.ResponderAgentID)
}
