package run

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/require"

	"github.com/tidbcloud/AgentTeam/internal/core/domain"
	"github.com/tidbcloud/AgentTeam/internal/workflow/plan"
)

func TestTaskWatchHandoffAllowsMoreThanLimitTokensAcrossDifferentTasks(t *testing.T) {
	t.Parallel()
	targetScope := taskWatchCardinalityScope("tsk_target")
	state, successor, predecessorIDs := taskWatchCardinalityState(
		domain.MaxWorkflowTaskWatchHandoffPredecessors+1,
		func(index int) domain.MessageScopeValue {
			if index == 0 {
				return targetScope
			}
			return taskWatchCardinalityScope(fmt.Sprintf("tsk_other_%03d", index))
		},
		func(index int) string { return fmt.Sprintf("agt_%03d", index) },
	)

	handoff, err := taskWatchHandoffForActivation(
		state,
		successor,
		"reviewer",
		nil,
		&targetScope,
		WorkScopeModeFixed,
		&targetScope,
	)
	require.NoError(t, err)
	require.Equal(t, []string{"agt_000"}, handoff.CausalPredecessorAgentIDs)
	require.NotContains(t, successor.TaskWatchFrontier, predecessorIDs[0])
	require.Len(t, successor.TaskWatchFrontier, len(predecessorIDs),
		"256 unrelated Task tokens plus the successor token remain a legal carrier")
}

func TestTaskWatchHandoffAllowsMoreThanLimitTokensForSameAgent(t *testing.T) {
	t.Parallel()
	targetScope := taskWatchCardinalityScope("tsk_target")
	state, successor, _ := taskWatchCardinalityState(
		domain.MaxWorkflowTaskWatchHandoffPredecessors+1,
		func(int) domain.MessageScopeValue { return targetScope },
		func(int) string { return "agt_same" },
	)

	handoff, err := taskWatchHandoffForActivation(
		state,
		successor,
		"reviewer",
		nil,
		&targetScope,
		WorkScopeModeFixed,
		&targetScope,
	)
	require.NoError(t, err)
	require.Equal(t, []string{"agt_same"}, handoff.CausalPredecessorAgentIDs)
	require.Equal(t, []string{successor.ID}, successor.TaskWatchFrontier)
}

func TestTaskWatchHandoffRejectsMoreThanLimitUniqueAgentsOnSameTask(t *testing.T) {
	t.Parallel()
	targetScope := taskWatchCardinalityScope("tsk_target")
	state, successor, predecessorIDs := taskWatchCardinalityState(
		domain.MaxWorkflowTaskWatchHandoffPredecessors+1,
		func(int) domain.MessageScopeValue { return targetScope },
		func(index int) string { return fmt.Sprintf("agt_%03d", index) },
	)

	handoff, err := taskWatchHandoffForActivation(
		state,
		successor,
		"reviewer",
		nil,
		&targetScope,
		WorkScopeModeFixed,
		&targetScope,
	)
	require.ErrorContains(t, err, "must not contain more than 256 predecessor Agent IDs")
	require.Nil(t, handoff)
	require.Equal(t, predecessorIDs, successor.TaskWatchFrontier,
		"a rejected handoff must restore the original carrier")
	require.Nil(t, successor.CallbackTaskWatch,
		"a rejected handoff must restore the original callback projection")
}

func TestNestedTaskWatchFrontierAboveTokenLimitAllowsBoundedUniqueHandoff(t *testing.T) {
	t.Parallel()
	const nestedTokenCount = 300
	targetScope := taskWatchCardinalityScope("tsk_nested")
	state, _, predecessorIDs := taskWatchCardinalityState(
		nestedTokenCount,
		func(int) domain.MessageScopeValue { return targetScope },
		func(index int) string {
			return fmt.Sprintf(
				"agt_%03d",
				index%domain.MaxWorkflowTaskWatchHandoffPredecessors,
			)
		},
	)
	// The helper's successor becomes the dormant outer join carrier. A separate
	// callback activation receives the collapsed frontier afterwards.
	joinIndex := len(state.Activations) - 1
	join := &state.Activations[joinIndex]
	join.ID = "act_outer_join"
	join.NodeID = "outer_join"
	join.TaskWatchFrontier = nil
	join.CallbackTaskWatch = nil
	successorIndex := len(state.Activations)
	state.Activations = append(state.Activations, Activation{
		ID: "act_successor", Sequence: uint64(successorIndex + 1), NodeID: "successor",
	})
	join = &state.Activations[joinIndex]
	successor := &state.Activations[successorIndex]

	items := func(offset, count int) []WorkItemState {
		result := make([]WorkItemState, count)
		for index := range result {
			result[index] = WorkItemState{
				ID:                fmt.Sprintf("wfi_%03d", offset+index),
				Key:               fmt.Sprintf("item_%03d", offset+index),
				TaskWatchFrontier: []string{predecessorIDs[offset+index]},
			}
		}
		return result
	}
	state.Scopes = []Scope{
		{
			ID: "scp_nested_a", Kind: plan.ArtifactScopeFanOut,
			ParentLineage: []ScopeRef{{ScopeID: "scp_outer", Kind: plan.ArtifactScopeParallel}},
			FanOut:        &FanOutState{Items: items(0, nestedTokenCount/2)},
		},
		{
			ID: "scp_nested_b", Kind: plan.ArtifactScopeFanOut,
			ParentLineage: []ScopeRef{{ScopeID: "scp_outer", Kind: plan.ArtifactScopeParallel}},
			FanOut:        &FanOutState{Items: items(nestedTokenCount/2, nestedTokenCount/2)},
		},
	}

	require.NoError(t, collapseScopeTaskWatchFrontiers(
		state,
		"scp_outer",
		join.ID,
		"",
	))
	require.Len(t, join.TaskWatchFrontier, nestedTokenCount)
	require.NoError(t, mergeTaskWatchFrontierIntoActivation(
		state,
		successor.ID,
		join.TaskWatchFrontier,
	))
	join.TaskWatchFrontier = nil // ordinary routing moves, rather than copies, the carrier
	require.Len(t, successor.TaskWatchFrontier, nestedTokenCount)

	handoff, err := taskWatchHandoffForActivation(
		state,
		successor,
		"reviewer",
		nil,
		&targetScope,
		WorkScopeModeFixed,
		&targetScope,
	)
	require.NoError(t, err)
	require.Len(
		t,
		handoff.CausalPredecessorAgentIDs,
		domain.MaxWorkflowTaskWatchHandoffPredecessors,
	)
}

func taskWatchCardinalityState(
	count int,
	taskScope func(int) domain.MessageScopeValue,
	agentID func(int) string,
) (*State, *Activation, []string) {
	state := &State{Activations: make([]Activation, count+1)}
	predecessorIDs := make([]string, count)
	for index := 0; index < count; index++ {
		activationID := fmt.Sprintf("act_predecessor_%03d", index)
		predecessorIDs[index] = activationID
		state.Activations[index] = Activation{
			ID: activationID, Sequence: uint64(index + 1), NodeID: fmt.Sprintf("callback_%03d", index),
			CallbackTaskWatch: &CallbackTaskWatchState{
				TaskScope: taskScope(index), ResponderRef: fmt.Sprintf("agent-%03d", index),
				ResponderAgentID: agentID(index),
			},
		}
	}
	state.Activations[count] = Activation{
		ID: "act_successor", Sequence: uint64(count + 1), NodeID: "successor",
		TaskWatchFrontier: cloneTaskWatchFrontier(predecessorIDs),
	}
	return state, &state.Activations[count], predecessorIDs
}

func taskWatchCardinalityScope(scopeID string) domain.MessageScopeValue {
	return domain.MessageScopeValue{Kind: domain.MessageScopeKindTask, ScopeID: scopeID}
}