package job

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestBuildComponentNode(t *testing.T) {
	tests := []struct {
		name          string
		componentNode string
		expected      string
	}{
		{
			name:     "default node group",
			expected: defaultComponentNode,
		},
		{
			name:          "component node from options",
			componentNode: "general",
			expected:      "general",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			job, err := newTestBuilder(t).Build(context.Background(), Options{
				TiDBPool:      "tidb-pool",
				KeyspaceName:  "ks",
				JobType:       TypeShared,
				JobName:       "tidb-worker-shared-1-0",
				ExecID:        "1:0",
				ComponentNode: tt.componentNode,
			})
			require.NoError(t, err)
			require.Equal(t, tt.expected, job.Spec.Template.Spec.NodeSelector[labels.ServerlessComponentNodeLabelKey])
			require.Equal(t, []corev1.Toleration{
				{
					Key:      labels.ServerlessComponentNodeLabelKey,
					Operator: corev1.TolerationOpEqual,
					Value:    tt.expected,
					Effect:   corev1.TaintEffectNoSchedule,
				},
			}, job.Spec.Template.Spec.Tolerations)
		})
	}
}

func newTestBuilder(t *testing.T) *Builder {
	t.Helper()

	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))

	workerConfigData := map[string]string{
		configFileKey:    "",
		startupScriptKey: "",
		types.TidbResourcesKey: `
requests:
  cpu: "1"
  memory: 1Gi
`,
	}

	client := fake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(
			&corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.ServerlessPoolConfig,
					Namespace: "tidb-pool",
				},
				Data: map[string]string{
					types.TidbImageKey: "tidb:test",
				},
			},
			&corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name:      scalerConfigKey,
					Namespace: "tidb-worker",
				},
				Data: workerConfigData,
			},
		).
		Build()

	return NewBuilder("tidb-worker", client)
}