package run

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

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

func countWorkItemStatus(scope *Scope, status WorkItemStatus) int {
	count := 0
	if scope == nil || scope.FanOut == nil {
		return count
	}
	for _, item := range scope.FanOut.Items {
		if item.Status == status {
			count++
		}
	}
	return count
}

func workItemRaw(key string, dependsOn []string, responder string) map[string]any {
	dependencies := make([]any, len(dependsOn))
	for index, dependency := range dependsOn {
		dependencies[index] = dependency
	}
	return map[string]any{
		plan.WorkItemFieldKey:           key,
		plan.WorkItemFieldDependsOn:     dependencies,
		plan.WorkItemFieldWorkScope:     map[string]any{"kind": "task", "scopeId": "tsk_" + key},
		plan.WorkItemFieldRelatedScopes: []any{map[string]any{"kind": "task", "scopeId": "tsk_parent"}},
		plan.WorkItemFieldResponder:     responder,
		plan.WorkItemFieldInstruction:   "implement " + key,
	}
}

func numericFanOutRuntimePlan() string {
	return strings.ReplaceAll(
		fanOutRuntimePlan,
		"answer: {type: string, required: true}",
		"answer: {type: number, required: true}",
	)
}

func fanOutSwitchChainRuntimePlan(count int) string {
	var nodes strings.Builder
	var edges strings.Builder
	for index := 0; index < count; index++ {
		fmt.Fprintf(&nodes, `  - id: route_%02d
    type: switch
    input:
      answer: {type: string, required: true}
    switch:
      mode: first_match
      cases:
      - port: empty
        when: 'input.answer == ""'
      - port: next
        isDefault: true
`, index)
		targetNode := "collected"
		targetPort := "item"
		if index+1 < count {
			targetNode = fmt.Sprintf("route_%02d", index+1)
			targetPort = "in"
		}
		fmt.Fprintf(&edges, `  - id: routed_%02d_next
    source: {nodeId: route_%02d, port: next}
    target: {nodeId: %s, port: %s}
    bindings:
      answer: {ref: input.answer}
  - id: routed_%02d_empty
    source: {nodeId: route_%02d, port: empty}
    target: {nodeId: %s, port: %s}
    bindings:
      answer: {ref: input.answer}
`, index, index, targetNode, targetPort, index, index, targetNode, targetPort)
	}
	result := strings.Replace(
		fanOutRuntimePlan,
		"  - id: collected\n",
		nodes.String()+"  - id: collected\n",
		1,
	)
	result = strings.Replace(
		result,
		"    target: {nodeId: collected, port: item}\n    bindings:\n      answer: {ref: output.answer}\n",
		"    target: {nodeId: route_00, port: in}\n    bindings:\n      answer: {ref: output.answer}\n",
		1,
	)
	result = strings.Replace(
		result,
		"  - id: joined\n",
		edges.String()+"  - id: joined\n",
		1,
	)
	return result
}

func fanOutPostCallbackRuntimePlan(chainLength int) string {
	result := fanOutSwitchChainRuntimePlan(chainLength)
	result = strings.Replace(
		result,
		`  - id: success
    type: end
    end: {outcome: succeeded}
`,
		`  - id: post
    type: send_message
    input:
      results:
        type: array
        required: true
        items: {type: object, required: true}
    sendMessage:
      target: "#all"
      text: post-process
      completion: callback
      callbackOutput:
        ack: {type: string, required: true}
  - id: success
    type: end
    end: {outcome: succeeded}
`,
		1,
	)
	result = strings.Replace(
		result,
		`    target: {nodeId: success, port: in}
    bindings:
      results: {ref: output.results}
`,
		`    target: {nodeId: post, port: in}
    bindings:
      results: {ref: output.results}
  - id: post_completed
    source: {nodeId: post, port: completed}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: input.results}
`,
		1,
	)
	return result
}

func fanOutCausalRepeatedBindingRuntimePlan() string {
	result := fanOutPostCallbackRuntimePlan(1)
	result = strings.Replace(
		result,
		`  - id: success
    type: end
    end: {outcome: succeeded}
`,
		`  - id: wide
    type: switch
    input:
      a: {type: string, required: true}
      b: {type: string, required: true}
      results:
        type: array
        required: true
        items: {type: object, required: true}
    switch:
      mode: first_match
      cases:
      - port: empty
        when: 'input.a == ""'
      - port: next
        isDefault: true
  - id: success
    type: end
    end: {outcome: succeeded}
`,
		1,
	)
	result = strings.Replace(
		result,
		`  - id: post_completed
    source: {nodeId: post, port: completed}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: input.results}
`,
		`  - id: post_completed
    source: {nodeId: post, port: completed}
    target: {nodeId: wide, port: in}
    bindings:
      a: {ref: output.ack}
      b: {ref: output.ack}
      results: {ref: input.results}
  - id: wide_next
    source: {nodeId: wide, port: next}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: input.results}
  - id: wide_empty
    source: {nodeId: wide, port: empty}
    target: {nodeId: success, port: in}
    bindings:
      results: {ref: input.results}
`,
		1,
	)
	return result
}

func fanOutRepeatedStaticJoinRuntimePlan(fieldCount int) string {
	var resultSchema strings.Builder
	var merge strings.Builder
	for index := 0; index < fieldCount; index++ {
		fmt.Fprintf(&resultSchema, `        result_%02d:
          type: array
          required: true
          items: {type: object, required: true}
`, index)
		fmt.Fprintf(&merge, "        result_%02d: {ref: items}\n", index)
	}
	result := strings.Replace(
		fanOutRuntimePlan,
		`      result:
        results:
          type: array
          required: true
          items: {type: object, required: true}
      merge:
        results: {ref: items}
`,
		"      result:\n"+resultSchema.String()+"      merge:\n"+merge.String(),
		1,
	)
	return strings.Replace(
		result,
		"      results: {ref: output.results}\n",
		"      results: {ref: output.result_00}\n",
		1,
	)
}

func startFanOut(t *testing.T, kernel *Kernel, workItems []any) *State {
	t.Helper()
	state, err := kernel.Start(fanOutPlanInput)
	require.NoError(t, err)
	planner, err := kernel.Step(state, reducerNow)
	require.NoError(t, err)
	require.Len(t, planner.Intents, 1)
	require.Equal(t, []string{"@dev1", "@dev2"}, planner.Intents[0].WorkItemsAllowedResponders)
	completed, err := kernel.Complete(planner.State, Completion{
		ActivationID: planner.Intents[0].ActivationID,
		Port:         "completed",
		Output: Values{
			"workScope": fanOutPlanInput["parent_task"],
			"workItems": workItems,
		},
	}, 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")
	return opened.State
}

func onlyFanOutScope(t *testing.T, state *State) *Scope {
	t.Helper()
	for index := range state.Scopes {
		if state.Scopes[index].FanOut != nil {
			return &state.Scopes[index]
		}
	}
	t.Fatal("missing fan-out scope")
	return nil
}

func itemStatuses(scope *Scope) []WorkItemStatus {
	result := make([]WorkItemStatus, len(scope.FanOut.Items))
	for index, item := range scope.FanOut.Items {
		result[index] = item.Status
	}
	return result
}

func stepToMessageIntent(t *testing.T, kernel *Kernel, state *State) (Intent, *State) {
	t.Helper()
	transition, err := kernel.Step(state, reducerNow)
	require.NoError(t, err)
	require.Len(t, transition.Intents, 1)
	require.Equal(t, IntentSendMessage, transition.Intents[0].Kind)
	return transition.Intents[0], transition.State
}

func completeIntent(
	t *testing.T,
	kernel *Kernel,
	state *State,
	intent Intent,
	port string,
	output Values,
	failure *Failure,
) *State {
	t.Helper()
	if port == "completed" && intent.WorkScope != nil {
		if output == nil {
			output = Values{}
		}
		if _, found := output["workScope"]; !found {
			output["workScope"] = messageScopeObject(*intent.WorkScope)
		}
	}
	transition, err := kernel.Complete(state, Completion{
		ActivationID: intent.ActivationID,
		Port:         port,
		Output:       output,
		Failure:      failure,
	}, reducerNow)
	require.NoError(t, err)
	return transition.State
}

func completeIntentAsResponder(
	t *testing.T,
	kernel *Kernel,
	state *State,
	intent Intent,
	responderAgentID string,
	output Values,
) *State {
	t.Helper()
	activationBefore, err := state.activation(intent.ActivationID)
	require.NoError(t, err)
	alreadyObserved := activationBefore.CallbackTaskWatch != nil &&
		activationBefore.CallbackTaskWatch.ResponderAgentID == responderAgentID
	if intent.WorkScope != nil {
		if output == nil {
			output = Values{}
		}
		output["workScope"] = messageScopeObject(*intent.WorkScope)
	}
	transition, err := kernel.Complete(state, Completion{
		ActivationID:             intent.ActivationID,
		Port:                     "completed",
		Output:                   output,
		CallbackResponderAgentID: responderAgentID,
	}, reducerNow)
	require.NoError(t, err)
	if alreadyObserved {
		require.NotContains(t, fanOutEventKinds(transition.Events), "fan_out_item_responder_activated")
	} else {
		require.Contains(t, fanOutEventKinds(transition.Events), "fan_out_item_responder_activated")
	}
	for _, event := range transition.Events {
		if event.Kind == "fan_out_item_responder_activated" {
			require.Equal(t, responderAgentID, event.AgentID)
		}
	}
	return transition.State
}

func fanOutEventKinds(events []Event) []string {
	kinds := make([]string, len(events))
	for index := range events {
		kinds[index] = events[index].Kind
	}
	return kinds
}

func activationItemKey(t *testing.T, state *State, activationID string) string {
	t.Helper()
	activation, err := state.activation(activationID)
	require.NoError(t, err)
	for _, reference := range activation.Lineage {
		if reference.Kind == plan.ArtifactScopeFanOut {
			return reference.ItemKey
		}
	}
	t.Fatalf("activation %s has no fan-out lineage", activationID)
	return ""
}

func activationForItem(t *testing.T, state *State, key string, status ActivationStatus) *Activation {
	t.Helper()
	for index := range state.Activations {
		activation := &state.Activations[index]
		if activation.Status != status {
			continue
		}
		for _, reference := range activation.Lineage {
			if reference.Kind == plan.ArtifactScopeFanOut && reference.ItemKey == key {
				return activation
			}
		}
	}
	t.Fatalf("missing %s activation for WorkItem %q; state=%s", status, key, formatStateForTest(state))
	return nil
}

func formatStateForTest(state *State) string {
	if state == nil {
		return "<nil>"
	}
	return fmt.Sprintf("%#v", state)
}