package replay

import (
	"testing"

	"github.com/stretchr/testify/require"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

	replayprovider "github.com/tidbcloud/eks-provider/controllers/replay/provider"
	"github.com/tidbcloud/eks-provider/controllers/replay/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/config"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/deploy"
)

func TestBuildReplayProviderSelectsTasksFromReplayContext(t *testing.T) {
	tests := []struct {
		name        string
		provider    string
		servicePlan string
		want        []string
		wantErr     string
	}{
		{
			name:        "eks premium target reuses existing workers",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypePremium,
			want:        []string{"Guardian", "StorageClass", "Namespace", "Service", "Replay"},
		},
		{
			name:        "ack premium target reuses existing workers",
			provider:    consts.ProviderTypeACK,
			servicePlan: consts.ServiceTypePremium,
			want:        []string{"Guardian", "StorageClass", "Namespace", "Service", "Replay"},
		},
		{
			name:        "ack premium target uses ack task plan",
			provider:    consts.ProviderTypeACK,
			servicePlan: consts.ServiceTypePremium,
			want:        []string{"Guardian", "StorageClass", "Namespace", "Service", "Replay"},
		},
		{
			name:        "eks dedicated uses dedicated nodepool path",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypeDedicated,
			want:        []string{"Guardian", "ProviderConfig", "NodeRole", "NodePool", "StorageClass", "Namespace", "Service", "Replay"},
		},
		{
			name:        "ack dedicated is unsupported",
			provider:    consts.ProviderTypeACK,
			servicePlan: consts.ServiceTypeDedicated,
			wantErr:     "unknown replay provider: ack",
		},
		{
			name:        "unknown provider returns unsupported",
			provider:    "unknown",
			servicePlan: consts.ServiceTypeDedicated,
			wantErr:     "unknown replay provider: unknown",
		},
		{
			name:        "unknown service plan falls back to dedicated provider",
			provider:    consts.ProviderTypeEKS,
			servicePlan: "unknown",
			want:        []string{"Guardian", "ProviderConfig", "NodeRole", "NodePool", "StorageClass", "Namespace", "Service", "Replay"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := newReplayTestContext(tt.provider, tt.servicePlan, cloudv1.Kubernetes{}, nil)
			provider, err := buildReplayProvider(c)
			if tt.wantErr != "" {
				require.EqualError(t, err, tt.wantErr)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tt.want, taskNames(provider.Tasks()))
		})
	}
}

func TestBuildReplayProviderRejectsNilReplay(t *testing.T) {
	tests := []struct {
		name        string
		provider    string
		servicePlan string
	}{
		{
			name:        "dedicated eks",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypeDedicated,
		},
		{
			name:        "premium eks",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypePremium,
		},
		{
			name:        "premium ack",
			provider:    consts.ProviderTypeACK,
			servicePlan: consts.ServiceTypePremium,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := &task.Context{
				RootConfig: config.RootConfig{
					ControllerConfig: config.ControllerConfig{
						Provider:    tt.provider,
						ServicePlan: tt.servicePlan,
					},
				},
			}
			got, err := buildReplayProvider(ctx)
			require.EqualError(t, err, "context replay is not initialized")
			require.Nil(t, got)
		})
	}
}

func TestDefaultReplayContextUsesProviderDefaults(t *testing.T) {
	tests := []struct {
		name         string
		provider     string
		servicePlan  string
		kubernetes   cloudv1.Kubernetes
		wantProvider any
		wantZones    []string
		wantSubnets  []string
		wantErr      string
	}{
		{
			name:         "eks uses node and secondary subnet groups",
			provider:     consts.ProviderTypeEKS,
			servicePlan:  consts.ServiceTypeDedicated,
			wantProvider: &replayprovider.DedicatedEKSReplayProvider{},
			kubernetes: cloudv1.Kubernetes{
				Cloud: cloudv1.Cloud{
					AWS: &cloudv1.CloudAWS{
						SubnetGroups: cloudv1.AWSSubnetGroups{
							{
								Name: cloudv1.AWSSubnetGroupNodes,
								Subnets: []cloudv1.AWSSubnet{
									{SubnetID: "subnet-node", Zone: "us-west-2a"},
								},
							},
							{
								Name: cloudv1.AWSSubnetGroupSecondaryCIDR,
								Subnets: []cloudv1.AWSSubnet{
									{SubnetID: "subnet-secondary", Zone: "us-west-2b"},
								},
							},
						},
					},
				},
			},
			wantZones:   []string{"us-west-2a", "us-west-2b"},
			wantSubnets: []string{"subnet-node", "subnet-secondary"},
		},
		{
			name:         "ack uses ack subnets",
			provider:     consts.ProviderTypeACK,
			servicePlan:  consts.ServiceTypePremium,
			wantProvider: &replayprovider.PremiumAliyunReplayProvider{},
			kubernetes: cloudv1.Kubernetes{
				Cloud: cloudv1.Cloud{
					Ack: &cloudv1.CloudAck{
						Subnets: []string{"vsw-1", "vsw-2"},
					},
				},
			},
			wantSubnets: []string{"vsw-1", "vsw-2"},
		},
		{
			name:         "eks premium reuses dataplane worker nodes",
			provider:     consts.ProviderTypeEKS,
			servicePlan:  consts.ServiceTypePremium,
			wantProvider: &replayprovider.PremiumEKSReplayProvider{},
			kubernetes: cloudv1.Kubernetes{
				Cloud: cloudv1.Cloud{
					AWS: &cloudv1.CloudAWS{
						SubnetGroups: cloudv1.AWSSubnetGroups{
							{
								Name: cloudv1.AWSSubnetGroupNodes,
								Subnets: []cloudv1.AWSSubnet{
									{SubnetID: "subnet-premium-node", Zone: "us-west-2c"},
								},
							},
							{
								Name: cloudv1.AWSSubnetGroupSecondaryCIDR,
								Subnets: []cloudv1.AWSSubnet{
									{SubnetID: "subnet-premium-secondary", Zone: "us-west-2d"},
								},
							},
						},
					},
				},
			},
			wantZones:   []string{"us-west-2c", "us-west-2d"},
			wantSubnets: []string{"subnet-premium-node", "subnet-premium-secondary"},
		},
		{
			name:        "ack dedicated is rejected",
			provider:    consts.ProviderTypeACK,
			servicePlan: consts.ServiceTypeDedicated,
			wantErr:     "unknown replay provider: ack",
		},
		{
			name:         "unknown service plan falls back to dedicated provider",
			provider:     consts.ProviderTypeEKS,
			servicePlan:  "unknown",
			wantProvider: &replayprovider.DedicatedEKSReplayProvider{},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := newReplayTestContext(tt.provider, tt.servicePlan, tt.kubernetes, []cloudv1.ReplayTask{{}})

			got, err := DefaultReplayContext(ctx)
			if tt.wantErr != "" {
				require.EqualError(t, err, tt.wantErr)
				return
			}
			require.NoError(t, err)
			require.NotNil(t, got.Provider)
			require.IsType(t, tt.wantProvider, got.Provider)
			require.ElementsMatch(t, tt.wantZones, got.Zones)
			require.ElementsMatch(t, tt.wantSubnets, got.Subnets)
		})
	}
}

func taskNames(tasks []deploy.Task) []string {
	names := make([]string, 0, len(tasks))
	for _, task := range tasks {
		names = append(names, task.Name())
	}
	return names
}

func TestBuildReplayProviderReturnsProviderSpecificAffinity(t *testing.T) {
	tests := []struct {
		name            string
		provider        string
		servicePlan     string
		replayTask      cloudv1.ReplayTask
		wantAffinity    []corev1.NodeSelectorRequirement
		wantTolerations []corev1.Toleration
	}{
		{
			name:        "eks premium defaults to node spec workers",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypePremium,
			replayTask:  cloudv1.ReplayTask{},
			wantAffinity: []corev1.NodeSelectorRequirement{
				{
					Key:      "node-spec",
					Operator: corev1.NodeSelectorOpIn,
					Values:   []string{task.DefaultReplayNodeSpec},
				},
			},
			wantTolerations: []corev1.Toleration{
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "worker",
					Effect:   corev1.TaintEffectNoSchedule,
				},
			},
		},
		{
			name:        "eks dedicated targets replay nodepool",
			provider:    consts.ProviderTypeEKS,
			servicePlan: consts.ServiceTypeDedicated,
			replayTask:  cloudv1.ReplayTask{},
			wantAffinity: []corev1.NodeSelectorRequirement{
				{
					Key:      task.NodeSelectorKey,
					Operator: corev1.NodeSelectorOpIn,
					Values:   []string{"r1"},
				},
				{
					Key:      task.NodeInstanceTypeKey,
					Operator: corev1.NodeSelectorOpIn,
					Values:   []string{task.DefaultInstanceType},
				},
			},
			wantTolerations: []corev1.Toleration{
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-replay",
					Effect:   corev1.TaintEffectNoSchedule,
				},
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := newReplayTestContext(tt.provider, tt.servicePlan, cloudv1.Kubernetes{}, []cloudv1.ReplayTask{{}})
			ctx.Replay.ObjectMeta = metav1.ObjectMeta{Name: "r1"}
			provider, err := buildReplayProvider(ctx)
			require.NoError(t, err)

			affinity := provider.ReplayNodeAffinity(tt.replayTask)
			require.NotNil(t, affinity)
			exprs := affinity.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions
			require.Equal(t, tt.wantAffinity, exprs)
			require.Equal(t, tt.wantTolerations, provider.ReplayTolerations())
		})
	}
}

func TestBuildReplayProviderUsesPerTaskInstanceType(t *testing.T) {
	ctx := newReplayTestContext(consts.ProviderTypeEKS, consts.ServiceTypePremium, cloudv1.Kubernetes{}, []cloudv1.ReplayTask{
		{
			Name:      "task-a",
			Resources: &cloudv1.ReplayTaskResources{InstanceType: "16C32G"},
		},
		{
			Name:      "task-b",
			Resources: &cloudv1.ReplayTaskResources{InstanceType: "32C64G"},
		},
	})

	provider, err := buildReplayProvider(ctx)
	require.NoError(t, err)

	affinityA := provider.ReplayNodeAffinity(ctx.Replay.Spec.Tasks[0])
	affinityB := provider.ReplayNodeAffinity(ctx.Replay.Spec.Tasks[1])
	require.NotNil(t, affinityA)
	require.NotNil(t, affinityB)

	exprsA := affinityA.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions
	exprsB := affinityB.NodeAffinity.RequiredDuringSchedulingIgnoredDuringExecution.NodeSelectorTerms[0].MatchExpressions
	require.Equal(t, []string{"16C32G"}, exprsA[0].Values)
	require.Equal(t, []string{"32C64G"}, exprsB[0].Values)
}

func TestDefaultReplayContextDoesNotMutatePremiumTaskResources(t *testing.T) {
	ctx := newReplayTestContext(consts.ProviderTypeEKS, consts.ServiceTypePremium, cloudv1.Kubernetes{}, []cloudv1.ReplayTask{{}})
	ctx.Replay.ObjectMeta = metav1.ObjectMeta{Name: "r1"}

	got, err := DefaultReplayContext(ctx)
	require.NoError(t, err)
	require.NotNil(t, got.Provider)
	require.Equal(t, []string{"Guardian", "StorageClass", "Namespace", "Service", "Replay"}, taskNames(got.Provider.Tasks()))
	require.Nil(t, got.Replay.Spec.Tasks[0].Resources)
}

func newReplayTestContext(provider, servicePlan string, kubernetes cloudv1.Kubernetes, tasks []cloudv1.ReplayTask) *task.Context {
	return &task.Context{
		Replay: &cloudv1.Replay{
			Spec: cloudv1.ReplaySpec{
				Kubernetes: kubernetes,
				Tasks:      tasks,
			},
		},
		RootConfig: config.RootConfig{
			ControllerConfig: config.ControllerConfig{
				Provider:    provider,
				ServicePlan: servicePlan,
			},
		},
	}
}