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

const fanOutRuntimePlan = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  parent_task:
    type: scope
    required: true
  dev1:
    type: agent
    required: true
  dev2:
    type: agent
    required: true
output:
  results:
    type: array
    required: true
    items: {type: object, required: true}
graph:
  entry:
    nodeId: planner
    port: in
    input:
      parent_task: {ref: plan.input.parent_task}
      dev1: {ref: plan.input.dev1}
      dev2: {ref: plan.input.dev2}
  nodes:
  - id: planner
    type: send_message
    input:
      parent_task: {type: scope, required: true}
      dev1: {type: agent, required: true}
      dev2: {type: agent, required: true}
    sendMessage:
      target: "#all"
      text: plan work
      completion: callback
      workScope: {ref: input.parent_task}
      workItems:
        allowedResponders: {expr: "[input.dev1, input.dev2]"}
  - id: dispatch
    type: fan_out
    input:
      work_items: {type: work_items, required: true}
    fanOut:
      items: {ref: input.work_items}
      maxConcurrency: 2
      failureMode: collect
  - id: implement
    type: send_message
    input:
      key: {type: string, required: true}
      depends_on:
        type: array
        required: true
        items: {type: string, required: true}
      work_scope: {type: scope, required: true}
      related_scopes:
        type: array
        required: true
        items: {type: scope, required: true}
      responder: {type: responder, required: true}
      instruction: {type: string, required: true}
      dependencies: {type: object, required: true}
    sendMessage:
      scope: {ref: input.work_scope}
      responder: {ref: input.responder}
      text: "implement {{ input.instruction }}"
      completion: callback
      workScope: {ref: input.work_scope}
      relatedScopes: {ref: input.related_scopes}
      callbackOutput:
        answer: {type: string, required: true}
  - id: collected
    type: join
    input:
      answer: {type: string, required: true}
    join:
      fanOutId: dispatch
      mode: all
      result:
        results:
          type: array
          required: true
          items: {type: object, required: true}
      merge:
        results: {ref: items}
  - id: success
    type: end
    end: {outcome: succeeded}
  edges:
  - id: planned
    source: {nodeId: planner, port: completed}
    target: {nodeId: dispatch, port: in}
    bindings:
      work_items: {ref: output.workItems}
  - id: dispatch_item
    source: {nodeId: dispatch, port: item}
    target: {nodeId: implement, port: in}
    bindings:
      key: {ref: item.key}
      depends_on: {ref: item.dependsOn}
      work_scope: {ref: item.workScope}
      related_scopes: {ref: item.relatedScopes}
      responder: {ref: item.responder}
      instruction: {ref: item.instruction}
      dependencies: {ref: item.dependencies}
  - id: implemented
    source: {nodeId: implement, port: completed}
    target: {nodeId: collected, port: item}
    bindings:
      answer: {ref: output.answer}
  - id: joined
    source: {nodeId: collected, port: completed}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: output.results}
`

const directFanOutRuntimePlan = `apiVersion: loop.plan/v1
kind: PlanVersion
input:
  parent_task: {type: scope, required: true}
  dev: {type: agent, required: true}
output:
  results:
    type: array
    required: true
    items: {type: object, required: true}
graph:
  entry:
    nodeId: planner
    port: in
    input:
      parent_task: {ref: plan.input.parent_task}
      dev: {ref: plan.input.dev}
  nodes:
  - id: planner
    type: send_message
    input:
      parent_task: {type: scope, required: true}
      dev: {type: agent, required: true}
    sendMessage:
      target: "#all"
      text: plan direct work
      completion: callback
      workScope: {ref: input.parent_task}
      workItems:
        allowedResponders: {expr: "[input.dev]"}
  - id: dispatch
    type: fan_out
    input:
      work_items: {type: work_items, required: true}
    fanOut:
      items: {ref: input.work_items}
      maxConcurrency: 1
      failureMode: collect
  - id: collected
    type: join
    join:
      fanOutId: dispatch
      mode: all
      result:
        results:
          type: array
          required: true
          items: {type: object, required: true}
      merge:
        results: {ref: items}
  - id: success
    type: end
    end: {outcome: succeeded}
  edges:
  - id: planned
    source: {nodeId: planner, port: completed}
    target: {nodeId: dispatch, port: in}
    bindings:
      work_items: {ref: output.workItems}
  - id: direct
    source: {nodeId: dispatch, port: item}
    target: {nodeId: collected, port: item}
  - id: joined
    source: {nodeId: collected, port: completed}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: output.results}
`

var fanOutPlanInput = Values{
	"parent_task": map[string]any{"kind": "task", "scopeId": "tsk_parent"},
	"dev1":        "@dev1",
	"dev2":        "@dev2",
}

func TestNormalizeWorkItemsNormalizedContractAndDAG(t *testing.T) {
	t.Parallel()
	raw := []any{
		workItemRaw("root", nil, "@dev1"),
		workItemRaw("child", []string{"root"}, "@dev2"),
	}
	delete(raw[0].(map[string]any), plan.WorkItemFieldDependsOn)
	items, err := NormalizeWorkItems("callback.workItems", raw, []string{" @dev1 "})
	require.ErrorContains(t, err, "surrounding whitespace")

	items, err = NormalizeWorkItems("callback.workItems", raw, []string{"@dev1", "@dev2", "@dev1"})
	require.NoError(t, err)
	require.Empty(t, items[0].DependsOn, "omitted dependsOn normalizedizes to []")
	require.Equal(t, []string{"root"}, items[1].DependsOn)
	normalized := WorkItemsValue(items)
	renormalized, err := NormalizeWorkItems("callback.workItems", normalized, nil)
	require.NoError(t, err)
	require.Equal(t, items, renormalized)

	cases := []struct {
		name string
		raw  []any
		want string
	}{
		{
			name: "unknown field",
			raw: []any{func() map[string]any {
				item := workItemRaw("a", nil, "@dev1")
				item["forged"] = true
				return item
			}()},
			want: "undeclared field",
		},
		{
			name: "unknown dependency",
			raw:  []any{workItemRaw("a", []string{"missing"}, "@dev1")},
			want: "unknown key",
		},
		{
			name: "cycle",
			raw: []any{
				workItemRaw("a", []string{"b"}, "@dev1"),
				workItemRaw("b", []string{"a"}, "@dev2"),
			},
			want: "dependency cycle",
		},
		{
			name: "responder allowlist",
			raw:  []any{workItemRaw("a", nil, "@other")},
			want: "not an allowed responder",
		},
	}
	for _, test := range cases {
		t.Run(test.name, func(t *testing.T) {
			_, err := NormalizeWorkItems("callback.workItems", test.raw, []string{"@dev1", "@dev2"})
			require.ErrorContains(t, err, test.want)
		})
	}

	tooMany := make([]any, plan.MaxWorkItems+1)
	for index := range tooMany {
		tooMany[index] = workItemRaw(fmt.Sprintf("item-%d", index), nil, "@dev1")
	}
	_, err = NormalizeWorkItems("callback.workItems", tooMany, []string{"@dev1"})
	require.ErrorContains(t, err, fmt.Sprintf("more than %d", plan.MaxWorkItems))

	tooLong := workItemRaw(string(make([]byte, plan.MaxWorkItemKeyBytes+1)), nil, "@dev1")
	_, err = NormalizeWorkItems("callback.workItems", []any{tooLong}, []string{"@dev1"})
	require.ErrorContains(t, err, fmt.Sprintf("exceed %d bytes", plan.MaxWorkItemKeyBytes))
}

func TestNormalizeWorkItemsAggregatePersistenceBudgets(t *testing.T) {
	t.Parallel()
	t.Run("maximum aggregate graph and related scopes fit normalized payload", func(t *testing.T) {
		raw := make([]any, plan.MaxWorkItems)
		keys := make([]string, len(raw))
		for index := range keys {
			keys[index] = fmt.Sprintf("item_%03d", index)
		}
		remainingEdges := plan.MaxWorkItemDependencyEdges
		for index := range raw {
			dependencyCount := index
			if dependencyCount > remainingEdges {
				dependencyCount = remainingEdges
			}
			dependencies := append([]string(nil), keys[:dependencyCount]...)
			remainingEdges -= dependencyCount
			item := workItemRaw(keys[index], dependencies, "@dev1")
			item[plan.WorkItemFieldRelatedScopes] = []any{
				map[string]any{"kind": "task", "scopeId": "tsk_parent"},
				map[string]any{"kind": "thread", "scopeId": "thr_context"},
			}
			raw[index] = item
		}
		require.Zero(t, remainingEdges)

		items, err := NormalizeWorkItems("callback.workItems", raw, []string{"@dev1"})
		require.NoError(t, err)
		payload, err := json.Marshal(map[string]any{"workItems": WorkItemsValue(items)})
		require.NoError(t, err)
		_, err = invocation.DeterministicJSONObject("workItems envelope", payload, 4<<20)
		require.NoError(t, err, "every maximum-budget WorkItems value must fit the ordinary persistence value-count boundary")
	})

	t.Run("dependency edge overflow is rejected before DAG persistence", func(t *testing.T) {
		raw := make([]any, 65)
		keys := make([]string, len(raw))
		for index := range keys {
			keys[index] = fmt.Sprintf("dense_%02d", index)
		}
		for index := range raw {
			raw[index] = workItemRaw(keys[index], keys[:index], "@dev1")
		}
		_, err := NormalizeWorkItems("callback.workItems", raw, []string{"@dev1"})
		require.ErrorContains(t, err, fmt.Sprintf(
			"more than %d dependency edges",
			plan.MaxWorkItemDependencyEdges,
		))
	})

	t.Run("aggregate related scope overflow is rejected", func(t *testing.T) {
		raw := make([]any, 9)
		for index := range raw {
			item := workItemRaw(fmt.Sprintf("context_%d", index), nil, "@dev1")
			related := make([]any, plan.MaxSendMessageRelatedScopes)
			for relatedIndex := range related {
				related[relatedIndex] = map[string]any{
					"kind":    "thread",
					"scopeId": fmt.Sprintf("thr_%d_%d", index, relatedIndex),
				}
			}
			item[plan.WorkItemFieldRelatedScopes] = related
			raw[index] = item
		}
		_, err := NormalizeWorkItems("callback.workItems", raw, []string{"@dev1"})
		require.ErrorContains(t, err, fmt.Sprintf(
			"more than %d related scope references",
			plan.MaxWorkItemRelatedScopeReferences,
		))
	})

	t.Run("aggregate normalized bytes are bounded", func(t *testing.T) {
		raw := make([]any, plan.MaxWorkItems)
		for index := range raw {
			item := workItemRaw(fmt.Sprintf("large_%03d", index), nil, "@dev1")
			item[plan.WorkItemFieldRelatedScopes] = []any{}
			item[plan.WorkItemFieldInstruction] = strings.Repeat("x", 5*1024)
			raw[index] = item
		}
		_, err := NormalizeWorkItems("callback.workItems", raw, []string{"@dev1"})
		require.ErrorContains(t, err, fmt.Sprintf(
			"must not exceed %d canonical JSON bytes",
			plan.MaxWorkItemsJSONBytes,
		))
	})
}

func TestFanOutZeroItemsCompletesOrderedJoin(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	state := startFanOut(t, kernel, nil)
	scope := onlyFanOutScope(t, state)
	require.Equal(t, ScopeSucceeded, scope.Status)
	require.Empty(t, scope.FanOut.Items)
	require.Equal(t, "completed", scope.FanOut.Outcome)

	terminal := stepUntilTerminal(t, kernel, state, 4)
	require.Equal(t, StatusSucceeded, terminal.Status)
	require.Equal(t, []any{}, terminal.Output["results"])
}

func TestFanOutPreservesNormalizedEmptyRelatedScopesFromPlannerWorkItems(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, fanOutRuntimePlan)
	item := workItemRaw("root", nil, "@dev1")
	item[plan.WorkItemFieldRelatedScopes] = []any{}

	state := startFanOut(t, kernel, []any{item})
	scope := onlyFanOutScope(t, state)
	require.Len(t, scope.FanOut.Items, 1)
	require.NotNil(t, scope.FanOut.Items[0].RelatedScopes)
	require.Empty(t, scope.FanOut.Items[0].RelatedScopes)
	_, err := kernel.Step(state, reducerNow)
	require.NoError(t, err, "normalized empty relatedScopes must pass the next reducer boundary")

	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	require.NotNil(t, onlyFanOutScope(t, restored).FanOut.Items[0].RelatedScopes)
}

func TestFanOutOneDirectItemPreservesEmptyTerminalValueAcrossRecovery(t *testing.T) {
	t.Parallel()
	kernel := compilePlanKernel(t, directFanOutRuntimePlan)
	input := Values{
		"parent_task": map[string]any{"kind": "task", "scopeId": "tsk_parent"},
		"dev":         "@dev1",
	}
	state, err := kernel.Start(input)
	require.NoError(t, err)
	planner, err := kernel.Step(state, reducerNow)
	require.NoError(t, err)
	require.Len(t, planner.Intents, 1)
	completed, err := kernel.Complete(planner.State, Completion{
		ActivationID: planner.Intents[0].ActivationID,
		Port:         "completed",
		Output: Values{
			"workScope": input["parent_task"],
			"workItems": []any{workItemRaw("only", nil, "@dev1")},
		},
	}, reducerNow)
	require.NoError(t, err)
	opened, err := kernel.Step(completed.State, reducerNow)
	require.NoError(t, err)
	require.Empty(t, opened.Intents, "fan-out open must not arm a join timeout intent")

	scope := onlyFanOutScope(t, opened.State)
	require.Equal(t, ScopeSucceeded, scope.Status)
	require.Equal(t, WorkItemSucceeded, scope.FanOut.Items[0].Status)
	require.NotNil(t, scope.FanOut.Items[0].Value)
	require.Empty(t, scope.FanOut.Items[0].Value)

	restored, err := JSONRoundTrip(opened.State)
	require.NoError(t, err)
	require.Equal(t, opened.State, restored)
	canceled, err := kernel.Cancel(restored)
	require.NoError(t, err)
	require.Equal(t, StatusCanceled, canceled.State.Status)
	replayedCancel, err := kernel.Cancel(canceled.State)
	require.NoError(t, err)
	require.Equal(t, canceled.State, replayedCancel.State)

	terminal := stepUntilTerminal(t, kernel, restored, 4)
	require.Equal(t, StatusSucceeded, terminal.Status)
	results := terminal.Output["results"].([]any)
	require.Len(t, results, 1)
	envelope := results[0].(map[string]any)
	require.Equal(t, map[string]any{}, envelope["value"])
	require.Equal(t, map[string]any{}, envelope[plan.FanOutItemFieldDependencies])
}

func TestFanOutConcurrencyCollectsFailedItemAndPreservesSourceOrder(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"),
	})
	scope := onlyFanOutScope(t, state)
	require.Equal(t, []WorkItemStatus{WorkItemRunning, WorkItemRunning, WorkItemPending}, itemStatuses(scope))

	aIntent, state := stepToMessageIntent(t, kernel, state)
	bIntent, state := stepToMessageIntent(t, kernel, state)
	require.NotEqual(t, aIntent.ActivationID, bIntent.ActivationID)
	state = completeIntent(t, kernel, state, aIntent, "completed", Values{"answer": "A"}, nil)
	scope = onlyFanOutScope(t, state)
	require.Equal(t, []WorkItemStatus{WorkItemSucceeded, WorkItemRunning, WorkItemRunning}, itemStatuses(scope))

	// Exact retry is idempotent and cannot launch another c lineage.
	replayed, err := kernel.Complete(state, Completion{
		ActivationID: aIntent.ActivationID,
		Port:         "completed",
		Output: Values{
			"answer":    "A",
			"workScope": messageScopeObject(*aIntent.WorkScope),
		},
	}, reducerNow)
	require.NoError(t, err)
	require.Equal(t, "completion_replayed", replayed.Events[0].Kind)
	require.Len(t, replayed.State.Activations, len(state.Activations))
	require.Equal(t, itemStatuses(onlyFanOutScope(t, state)), itemStatuses(onlyFanOutScope(t, replayed.State)))

	state = completeIntent(t, kernel, state, bIntent, "failed", Values{}, &Failure{Code: "agent.declined"})
	scope = onlyFanOutScope(t, state)
	require.Equal(t, WorkItemFailed, scope.FanOut.Items[1].Status)
	cIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntent(t, kernel, state, cIntent, "completed", Values{"answer": "C"}, nil)
	scope = onlyFanOutScope(t, state)
	require.Equal(t, ScopeSucceeded, scope.Status)

	terminal := stepUntilTerminal(t, kernel, state, 4)
	results := terminal.Output["results"].([]any)
	require.Len(t, results, 3)
	require.Equal(t, "a", results[0].(map[string]any)["key"])
	require.Equal(t, "b", results[1].(map[string]any)["key"])
	require.Equal(t, "c", results[2].(map[string]any)["key"])
	require.Equal(t, string(WorkItemFailed), results[1].(map[string]any)["status"])
}

func TestFanOutSuccessOnlyDAGSchedulingDependenciesAndTransitiveBlocking(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"),
	})
	scope := onlyFanOutScope(t, state)
	require.Equal(
		t,
		[]WorkItemStatus{WorkItemRunning, WorkItemPending, WorkItemPending, WorkItemRunning},
		itemStatuses(scope),
	)
	rootIntent, state := stepToMessageIntent(t, kernel, state)
	independentIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, "root", activationItemKey(t, state, rootIntent.ActivationID))
	require.Equal(t, "independent", activationItemKey(t, state, independentIntent.ActivationID))

	state = completeIntent(t, kernel, state, rootIntent, "completed", Values{"answer": "root result"}, nil)
	scope = onlyFanOutScope(t, state)
	require.Equal(t, WorkItemRunning, scope.FanOut.Items[1].Status)
	childActivation := activationForItem(t, state, "child", ActivationReady)
	dependencies := childActivation.Input["dependencies"].(map[string]any)
	require.Contains(t, dependencies, "root")
	rootEnvelope := dependencies["root"].(map[string]any)
	require.Equal(t, string(WorkItemSucceeded), rootEnvelope["status"])
	require.Equal(t, map[string]any{"answer": "root result"}, rootEnvelope["value"])
	require.NotContains(t, rootEnvelope, "output", "dependency envelopes expose the closing join input as value")
	require.NotContains(t, rootEnvelope, plan.FanOutItemFieldDependencies)

	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	require.Equal(t, state, restored)
	state = restored
	childIntent, state := stepToMessageIntent(t, kernel, state)
	state = completeIntent(t, kernel, state, childIntent, "failed", Values{}, &Failure{Code: "implementation.failed"})
	scope = onlyFanOutScope(t, state)
	require.Equal(t, WorkItemFailed, scope.FanOut.Items[1].Status)
	require.Equal(t, WorkItemBlocked, scope.FanOut.Items[2].Status)
	require.Equal(t, []string{"child"}, scope.FanOut.Items[2].BlockedBy)
	require.Equal(t, WorkItemRunning, scope.FanOut.Items[3].Status, "independent branch continues")

	state = completeIntent(t, kernel, state, independentIntent, "completed", Values{"answer": "independent"}, nil)
	scope = onlyFanOutScope(t, state)
	require.Equal(t, ScopeSucceeded, scope.Status)
	terminal := stepUntilTerminal(t, kernel, state, 4)
	results := terminal.Output["results"].([]any)
	require.Equal(t, string(WorkItemBlocked), results[2].(map[string]any)["status"])
	require.Equal(t, []any{"child"}, results[2].(map[string]any)["blockedBy"])
}

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

	rootIntent, state := stepToMessageIntent(t, kernel, state)
	require.Equal(t, "root", activationItemKey(t, state, rootIntent.ActivationID))
	state = completeIntent(
		t,
		kernel,
		state,
		rootIntent,
		"completed",
		Values{"answer": json.Number("42")},
		nil,
	)

	scope := onlyFanOutScope(t, state)
	require.Equal(t, WorkItemRunning, fanOutItemByKey(scope.FanOut, "child").Status)
	require.Equal(
		t,
		"add $root.answer and $$root.answer",
		fanOutItemByKey(scope.FanOut, "child").Instruction,
		"the frozen WorkItem keeps the authored instruction for audit",
	)
	childActivation := activationForItem(t, state, "child", ActivationReady)
	require.Equal(
		t,
		"add 42 and $root.answer",
		childActivation.Input[plan.WorkItemFieldInstruction],
		"the child activation freezes the single-pass materialized bytes",
	)
	dependencies := childActivation.Input[plan.FanOutItemFieldDependencies].(map[string]any)
	require.Equal(
		t,
		map[string]any{"answer": json.Number("42")},
		dependencies["root"].(map[string]any)["value"],
		"materialization does not remove or rewrite structured dependencies",
	)
	encodedItem, err := json.Marshal(fanOutItemByKey(scope.FanOut, "child"))
	require.NoError(t, err)
	require.NotContains(
		t,
		string(encodedItem),
		`"dependencies"`,
		"dependency results live once on predecessor Values and once in the launched child input, not on every WorkItem edge",
	)

	restored, err := JSONRoundTrip(state)
	require.NoError(t, err)
	restoredChild := activationForItem(t, restored, "child", ActivationReady)
	require.Equal(t, "add 42 and $root.answer", restoredChild.Input[plan.WorkItemFieldInstruction])
	childIntent, _ := stepToMessageIntent(t, kernel, restored)
	require.Equal(t, "implement add 42 and $root.answer", childIntent.Text)
}

func TestPlannerIntentFreezesFanOutDependencyValueContract(t *testing.T) {
	t.Parallel()
	t.Run("typed item value", func(t *testing.T) {
		kernel := compilePlanKernel(t, numericFanOutRuntimePlan())
		state, err := kernel.Start(fanOutPlanInput)
		require.NoError(t, err)
		transition, err := kernel.Step(state, reducerNow)
		require.NoError(t, err)
		require.Len(t, transition.Intents, 1)
		contract := transition.Intents[0].WorkItemsDependencyValueContract
		require.NotNil(t, contract)
		require.Equal(t, plan.ArtifactContract{{
			Name: "answer",
			Schema: plan.ArtifactSchema{
				Type: plan.TypeNumber, Required: true, Values: []string{},
			},
		}}, *contract)
	})

	t.Run("empty item value remains a present empty contract", func(t *testing.T) {
		kernel := compilePlanKernel(t, directFanOutRuntimePlan)
		state, err := kernel.Start(Values{
			"parent_task": fanOutPlanInput["parent_task"],
			"dev":         "@dev1",
		})
		require.NoError(t, err)
		transition, err := kernel.Step(state, reducerNow)
		require.NoError(t, err)
		require.Len(t, transition.Intents, 1)
		contract := transition.Intents[0].WorkItemsDependencyValueContract
		require.NotNil(t, contract)
		require.NotNil(t, *contract)
		require.Empty(t, *contract)
	})
}