package workflow

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	"github.com/tidbcloud/eks-provider/controllers/workflow/action"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
)

func TestSequenceScheduler_ScheduleOnRunning(t *testing.T) {
	ctx := context.Background()

	tests := map[string]struct {
		workflow      *cloudv1.Workflow
		wantScheduled bool
		wantErr       bool
		verify        func(t *testing.T, wf *cloudv1.Workflow)
	}{
		"skip scheduling when there are running steps": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:        cloudv1.WorkflowPhaseRunning,
					RunningSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseRunning},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"skip scheduling when there are ending steps": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:       cloudv1.WorkflowPhaseRunning,
					EndingSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseEnding},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"schedule first pending step to run": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
						{Name: "step2", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseRunning,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhasePending},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhasePending},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.RunningSteps, "step1")
				assert.Equal(t, cloudv1.WorkflowStepPhaseRunning, wf.MustStepStatus(ctx, "step1").Phase)
			},
		},
		"no scheduling when all steps are completed": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
						{Name: "step2", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseRunning,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCompleted},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhaseCompleted},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"schedule second step when first is succeeded": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
						{Name: "step2", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseRunning,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCompleted},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhasePending},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.RunningSteps, "step2")
				assert.Equal(t, cloudv1.WorkflowStepPhaseRunning, wf.MustStepStatus(ctx, "step2").Phase)
			},
		},
	}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			scheduler := NewSequenceScheduler(action.DefaultRegistry)
			scheduled, err := scheduler.scheduleOnRunning(ctx, tt.workflow)

			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}

			assert.Equal(t, tt.wantScheduled, scheduled)

			if tt.verify != nil {
				tt.verify(t, tt.workflow)
			}
		})
	}
}

func TestSequenceScheduler_ScheduleOnCancelling(t *testing.T) {
	ctx := context.Background()

	type mockAction struct {
		action.Action
	}

	tests := map[string]struct {
		workflow      *cloudv1.Workflow
		wantScheduled bool
		wantErr       bool
		verify        func(t *testing.T, wf *cloudv1.Workflow)
	}{
		"skip scheduling when there are cancelling steps": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: consts.WorkflowActionCancelableSleep},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:           cloudv1.WorkflowPhaseCancelling,
					CancellingSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCancelling},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"skip scheduling when there are cancelEnding steps": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:             cloudv1.WorkflowPhaseCancelling,
					CancelEndingSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCancelEnding},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"schedule running step to cancel when action implements ActionCancel": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: consts.WorkflowActionCancelableSleep},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:        cloudv1.WorkflowPhaseCancelling,
					RunningSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseRunning},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.CancellingSteps, "step1")
				assert.Equal(t, cloudv1.WorkflowStepPhaseCancelling, wf.MustStepStatus(ctx, "step1").Phase)
			},
		},
		"keep running step going when action doesn't implement ActionCancel": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:        cloudv1.WorkflowPhaseCancelling,
					RunningSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseRunning},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"schedule ending step to cancel when action implements ActionCancel": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: consts.WorkflowActionCancelableSleep},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:       cloudv1.WorkflowPhaseCancelling,
					EndingSteps: []string{"step1"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseEnding},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.CancellingSteps, "step1")
				assert.Equal(t, cloudv1.WorkflowStepPhaseCancelling, wf.MustStepStatus(ctx, "step1").Phase)
			},
		},
		"schedule pending step to cancel": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
						{Name: "step2", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseCancelling,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhasePending},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhasePending},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.FinishedSteps, "step1")
				assert.Equal(t, cloudv1.WorkflowStepPhaseCancelled, wf.MustStepStatus(ctx, "step1").Phase)
			},
		},
		"keep running step when there are finished steps in same group": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action", StepGroup: "group1"},
						{Name: "step2", Uses: "test-action", StepGroup: "group1"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase:        cloudv1.WorkflowPhaseCancelling,
					RunningSteps: []string{"step2"},
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCompleted},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhaseRunning},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
		"schedule pending step to run when there are finished steps in same group": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action", StepGroup: "group1"},
						{Name: "step2", Uses: "test-action", StepGroup: "group1"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseCancelling,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCompleted},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhasePending},
					},
				},
			},
			wantScheduled: true,
			wantErr:       false,
			verify: func(t *testing.T, wf *cloudv1.Workflow) {
				assert.Contains(t, wf.Status.RunningSteps, "step2")
				assert.Equal(t, cloudv1.WorkflowStepPhaseRunning, wf.MustStepStatus(ctx, "step2").Phase)
			},
		},
		"no scheduling when all steps are completed": {
			workflow: &cloudv1.Workflow{
				ObjectMeta: metav1.ObjectMeta{
					Name: "test-workflow",
				},
				Spec: cloudv1.WorkflowSpec{
					Steps: []cloudv1.WorkflowStep{
						{Name: "step1", Uses: "test-action"},
						{Name: "step2", Uses: "test-action"},
					},
				},
				Status: cloudv1.WorkflowStatus{
					Phase: cloudv1.WorkflowPhaseCancelling,
					Steps: []*cloudv1.WorkflowStepStatus{
						{Name: "step1", Phase: cloudv1.WorkflowStepPhaseCompleted},
						{Name: "step2", Phase: cloudv1.WorkflowStepPhaseCancelled},
					},
				},
			},
			wantScheduled: false,
			wantErr:       false,
		},
	}

	r := &action.Registry{Actions: map[cloudv1.WorkflowActionName]action.Action{}}
	for name, act := range action.DefaultRegistry.Actions {
		r.Actions[name] = act
	}
	r.Actions["test-action"] = &mockAction{}

	for name, tt := range tests {
		t.Run(name, func(t *testing.T) {
			scheduler := NewSequenceScheduler(r)
			scheduled, err := scheduler.scheduleOnCancelling(ctx, tt.workflow)

			if tt.wantErr {
				assert.Error(t, err)
			} else {
				assert.NoError(t, err)
			}

			assert.Equal(t, tt.wantScheduled, scheduled)

			if tt.verify != nil {
				tt.verify(t, tt.workflow)
			}
		})
	}
}