package common

import (
	"fmt"
	"testing"

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

	"github.com/tidbcloud/eks-provider/controllers/cluster/task"
	"github.com/tidbcloud/eks-provider/controllers/cluster/task/common/feature"
	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/namer"
)

type fakeCloudConfigProvider struct {
	instanceTypes map[string]cloudv1.InstanceType
}

func (p fakeCloudConfigProvider) GetProviderType() string {
	return ""
}

func (p fakeCloudConfigProvider) GetInstanceTypes() ([]cloudv1.InstanceType, error) {
	instanceTypes := make([]cloudv1.InstanceType, 0, len(p.instanceTypes))
	for _, it := range p.instanceTypes {
		instanceTypes = append(instanceTypes, it)
	}
	return instanceTypes, nil
}

func (p fakeCloudConfigProvider) GetStorageClasses() ([]*storagev1.StorageClass, error) {
	return nil, nil
}

func (p fakeCloudConfigProvider) FindInstanceType(name string) (cloudv1.InstanceType, error) {
	if it, ok := p.instanceTypes[name]; ok {
		return it, nil
	}
	return cloudv1.InstanceType{}, fmt.Errorf("instance type %s not found", name)
}

func (p fakeCloudConfigProvider) FindRegionZones(_ string) ([]string, error) {
	return nil, nil
}

func TestTiDBClusterV2_buildNodeSelectorAndTolerations(t1 *testing.T) {
	tests := []struct {
		name         string
		tsk          *TiDBClusterV2
		comp         cloudv1.ComponentType
		instanceType string
		config       config.RootConfig
		provider     config.CloudConfigProvider
		wantNs       map[string]string
		wantTo       []corev1.Toleration
	}{
		{
			name:         "aws",
			tsk:          &TiDBClusterV2{},
			comp:         cloudv1.ComponentPD,
			instanceType: "t3.medium",
			wantNs: map[string]string{
				"gardener.wg.prefix.name":          "123456789-pd",
				"tidb-component":                   "pd",
				"tidb-namespace":                   "tidb123456789",
				"tidbcloud-cluster":                "123456789",
				"node.kubernetes.io/instance-type": "t3.medium",
			},
			wantTo: []corev1.Toleration{
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-pd",
					Effect:   corev1.TaintEffectNoSchedule,
				},
			},
		},
		{
			name: "gcp",
			tsk: &TiDBClusterV2{
				k8sVersion: "1.31",
				featureGates: feature.NewGates(
					feature.UseClusterAutoscaler,
					feature.AddK8sVersionToNodeSelector,
					feature.AddNodePoolDataPlaneLabel,
				),
			},
			comp:         cloudv1.ComponentPD,
			instanceType: "n1-standard-4",
			wantNs: map[string]string{
				"cloud.pingcap.com/node-group-type":   "data-plane",
				"node.kubernetes.io/instance-type":    "n1-standard-4",
				"tidbcloud.com/node-pool-k8s-version": "1.31",
			},
			wantTo: []corev1.Toleration{
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-pd",
					Effect:   corev1.TaintEffectNoSchedule,
				},
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-shared",
					Effect:   corev1.TaintEffectNoSchedule,
				},
			},
		},
		{
			name: "gcp arm",
			tsk: &TiDBClusterV2{
				k8sVersion: "1.35",
				featureGates: feature.NewGates(
					feature.UseClusterAutoscaler,
					feature.AddK8sVersionToNodeSelector,
					feature.AddNodePoolDataPlaneLabel,
					feature.Add3LayerNodeSelectorAndTolerations,
				),
			},
			comp:         cloudv1.ComponentPD,
			instanceType: "c4a-standard-2",
			config: config.RootConfig{
				ControllerConfig: config.ControllerConfig{
					Provider: consts.ProviderTypeGKE,
				},
			},
			provider: fakeCloudConfigProvider{
				instanceTypes: map[string]cloudv1.InstanceType{
					"c4a-standard-2": {
						Name: "c4a-standard-2",
						Arch: consts.ARM64ArchName,
					},
				},
			},
			wantNs: map[string]string{
				"cloud.pingcap.com/for-3-layer-arch":  "true",
				"cloud.pingcap.com/node-group-type":   "data-plane",
				"node.kubernetes.io/instance-type":    "c4a-standard-2",
				"tidbcloud.com/node-pool-k8s-version": "1.35",
			},
			wantTo: []corev1.Toleration{
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-pd",
					Effect:   corev1.TaintEffectNoSchedule,
				},
				{
					Key:      "dedicated",
					Operator: corev1.TolerationOpEqual,
					Value:    "tidb-shared",
					Effect:   corev1.TaintEffectNoSchedule,
				},
				{
					Key:      "3-layer-arch",
					Operator: corev1.TolerationOpEqual,
					Value:    "true",
					Effect:   corev1.TaintEffectNoSchedule,
				},
				{
					Key:      "kubernetes.io/arch",
					Operator: corev1.TolerationOpEqual,
					Value:    "arm64",
					Effect:   corev1.TaintEffectNoSchedule,
				},
			},
		},
	}

	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name: "123456789",
		},
	}
	baseTask := &task.BaseTask{
		Context: &task.Context{
			Namer: namer.NewNamer().
				WithClusterName(cluster.Name).
				WithTiDBNamespace(cluster.PrefixedId()),
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			bt := *baseTask
			ctx := *baseTask.Context
			ctx.Config = tt.config
			ctx.Provider = tt.provider
			bt.Context = &ctx
			tt.tsk.BaseTask = &bt
			if tt.tsk.featureGates == nil {
				tt.tsk.featureGates = feature.NewGates()
			}
			gotNs, gotTo := tt.tsk.buildNodeSelectorAndTolerations(tt.comp, tt.instanceType)
			assert.Equal(t1, tt.wantNs, gotNs)
			assert.Equal(t1, tt.wantTo, gotTo)
		})
	}
}