package mixeddeploypd

import (
	"context"
	"fmt"
	"testing"
	"time"

	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"sigs.k8s.io/controller-runtime/pkg/client"

	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/controllers/test"
	"github.com/tidbcloud/infra-api/controllers/utils"
)

func TestMixedDeployEnableChecking(t *testing.T) {
	type Case struct {
		name        string
		pd          int
		tikv        int
		tidb        int
		tiflash     int
		createTime  time.Time
		enabled     bool
		allows      []string
		blocks      []string
		namespace   string
		affectAfter time.Time
		result      bool
		tikvCores   int
	}
	now := time.Now()
	cases := []Case{
		{
			name:    "disabled",
			pd:      1,
			tikv:    1,
			tidb:    1,
			tiflash: 0,
			enabled: false,
			allows:  []string{"*"},
			result:  false,
		},
		{
			name:    "enabled",
			pd:      1,
			tikv:    1,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  true,
		},
		{
			name:        "affect after",
			pd:          1,
			tikv:        1,
			tidb:        1,
			tiflash:     0,
			createTime:  now.Add(time.Second),
			enabled:     true,
			allows:      []string{"*"},
			affectAfter: now,
			result:      true,
		},
		{
			name:        "affect creating",
			pd:          1,
			tikv:        1,
			tidb:        1,
			tiflash:     0,
			createTime:  now,
			enabled:     true,
			allows:      []string{"*"},
			affectAfter: now,
			result:      true,
		},
		{
			name:    "empty allow list",
			pd:      1,
			tikv:    1,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{},
			result:  false,
		},
		{
			name:    "not in allow list",
			pd:      1,
			tikv:    1,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"allow-namespace-1", "allow-namespace-2"},
			result:  false,
		},
		{
			name:      "in block list",
			pd:        1,
			tikv:      1,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"*"},
			blocks:    []string{"blocked-namespace"},
			namespace: "blocked-namespace",
			result:    false,
		},
		{
			name:      "in allow list",
			pd:        1,
			tikv:      1,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"allow-namespace-1", "allow-namespace", "allow-namespace-2"},
			blocks:    []string{"blocked-namespace"},
			namespace: "allow-namespace",
			result:    true,
		},
		{
			name:      "in both allow and block list",
			pd:        1,
			tikv:      1,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"allow-namespace-1", "cluster-namespace", "allow-namespace-2"},
			blocks:    []string{"cluster-namespace"},
			namespace: "cluster-namespace",
			result:    false,
		},
		{
			name:        "not affect before",
			pd:          1,
			tikv:        1,
			tidb:        1,
			tiflash:     0,
			createTime:  now,
			enabled:     true,
			allows:      []string{"*"},
			affectAfter: now.Add(time.Second),
			result:      false,
		},
		{
			name:    "has tiflash",
			pd:      1,
			tikv:    1,
			tidb:    1,
			tiflash: 1,
			enabled: true,
			allows:  []string{"*"},
			result:  false,
		},
		{
			name:    "2 tidb",
			pd:      3,
			tikv:    3,
			tidb:    2,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  true,
		},
		{
			name:    "more than 2 tidb",
			pd:      3,
			tikv:    3,
			tidb:    3,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  false,
		},
		{
			name:    "pd less than tikv",
			pd:      3,
			tikv:    1,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  false,
		},
		{
			name:    "3 tikv",
			pd:      1,
			tikv:    3,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  true,
		},
		{
			name:    "3 pd, 3 tikv",
			pd:      3,
			tikv:    3,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  true,
		},
		{
			name:    "more than 3 tikv",
			pd:      3,
			tikv:    4,
			tidb:    1,
			tiflash: 0,
			enabled: true,
			allows:  []string{"*"},
			result:  false,
		},
		{
			name:      "2c tikv",
			pd:        3,
			tikv:      3,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"*"},
			tikvCores: 2,
			result:    true,
		},
		{
			name:      "4c tikv",
			pd:        3,
			tikv:      3,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"*"},
			tikvCores: 4,
			result:    true,
		},
		{
			name:      "8c tikv",
			pd:        3,
			tikv:      3,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"*"},
			tikvCores: 8,
			result:    false,
		},
		{
			name:      "16c tikv",
			pd:        3,
			tikv:      3,
			tidb:      1,
			tiflash:   0,
			enabled:   true,
			allows:    []string{"*"},
			tikvCores: 16,
			result:    false,
		},
	}
	for _, it := range cases {
		pdMixedDeployConfig := &cloudv1.MixedDeployPDWithTiKV{
			Enabled:                   it.enabled,
			AffectClusterCreatedAfter: metav1.Time{Time: it.affectAfter},
			AffectNamespaceAllowList:  it.allows,
			AffectNamespaceBlockList:  it.blocks,
		}
		pdMixedDeployConfig.Init()
		if it.tikvCores == 0 {
			it.tikvCores = 4
		}
		tikvInstanceSpec := cloudv1.InstanceType{
			Name: "instance",
			CPU:  resource.MustParse(fmt.Sprintf("%d", it.tikvCores)),
		}
		result := shouldMixedDeployPDWithTiKV(it.pd, it.tikv, it.tidb, it.tiflash, "testing",
			it.namespace, "fake-cluster-name", tikvInstanceSpec, it.createTime, pdMixedDeployConfig)
		if result != it.result {
			_, reason := pdMixedDeployConfig.ClusterShouldBeAffected(it.pd, it.tikv, it.tidb, it.tiflash, tikvInstanceSpec, it.namespace, it.createTime)
			t.Fatal("test " + it.name + ", return reason: " + reason)
		}
	}
}

func TestMixedDeployNodesAllocChecking(t *testing.T) {
	type Case struct {
		name           string
		enabled        bool
		tcIsMixed      bool
		clusterPhase   cloudv1.ClusterPhase
		pdReplicas     int
		runningMixedPD int
		result         bool
	}

	cases := []Case{
		{
			name:         "disabled",
			enabled:      false,
			tcIsMixed:    true,
			clusterPhase: cloudv1.ClusterPhaseCreating,
			result:       false,
		},
		{
			name:           "tc is created before not-affected then become affected",
			enabled:        true,
			tcIsMixed:      false,
			clusterPhase:   cloudv1.ClusterPhaseCreating,
			pdReplicas:     3,
			runningMixedPD: 0,
			result:         false,
		},
		{
			name:           "phase creating start",
			enabled:        true,
			tcIsMixed:      true,
			clusterPhase:   cloudv1.ClusterPhaseCreating,
			pdReplicas:     3,
			runningMixedPD: 0,
			result:         true,
		},
		{
			name:           "phase creating finish",
			enabled:        true,
			tcIsMixed:      true,
			clusterPhase:   cloudv1.ClusterPhaseCreating,
			pdReplicas:     3,
			runningMixedPD: 3,
			result:         true,
		},
		{
			name:           "phase scale-in start",
			enabled:        true,
			tcIsMixed:      true,
			clusterPhase:   cloudv1.ClusterPhaseScaling,
			pdReplicas:     3,
			runningMixedPD: 0,
			result:         false,
		},
		{
			name:           "phase scale-in finish",
			enabled:        true,
			tcIsMixed:      true,
			clusterPhase:   cloudv1.ClusterPhaseScaling,
			pdReplicas:     3,
			runningMixedPD: 3,
			result:         true,
		},
	}

	now := time.Now()

	for i, it := range cases {
		mockPDPodList := func(_ context.Context, list client.ObjectList, _ ...client.ListOption) error {
			podList, isPodList := list.(*corev1.PodList)
			if isPodList {
				for j := 0; j < it.runningMixedPD; j++ {
					pod := corev1.Pod{
						Status: corev1.PodStatus{
							Phase: corev1.PodRunning,
						},
						Spec: corev1.PodSpec{
							Tolerations: []corev1.Toleration{
								{
									Key:      "dedicated",
									Operator: corev1.TolerationOpEqual,
									Effect:   corev1.TaintEffectNoSchedule,
									Value:    fmt.Sprintf("tidb-%s", string(cloudv1.ComponentTiKV)),
								},
							},
						},
					}
					podList.Items = append(podList.Items, pod)
				}
			}
			nodeList, isNodeList := list.(*corev1.NodeList)
			if isNodeList {
				if !it.result {
					for i := 0; i < it.pdReplicas; i++ {
						nodeList.Items = append(nodeList.Items, corev1.Node{})
					}
				}
			}
			return nil
		}
		mockShootClient := &test.MockClient{
			MockList: mockPDPodList,
		}
		tc := &v1alpha1.TidbCluster{
			ObjectMeta: metav1.ObjectMeta{
				CreationTimestamp: metav1.Time{Time: now},
			},
			Spec: v1alpha1.TidbClusterSpec{
				PD: &v1alpha1.PDSpec{
					ComponentSpec: v1alpha1.ComponentSpec{
						Tolerations: utils.Toleration(cloudv1.ComponentPD),
					},
				},
			},
		}
		tc.Namespace = "fake-ns"
		tc.Name = "fake-name"
		if it.tcIsMixed {
			tc.Spec.PD.Tolerations = utils.Toleration(cloudv1.ComponentTiKV)
		}
		cluster := &cloudv1.Cluster{
			Status: cloudv1.ClusterStatus{
				Phase: it.clusterPhase,
			},
			Spec: cloudv1.ClusterSpec{
				Cluster: cloudv1.ClusterTopology{
					TiKV: cloudv1.TiKV{
						InstanceType: "tikv-instance-type",
					},
				},
			},
		}
		pdMixedDeployConfig := &cloudv1.MixedDeployPDWithTiKV{
			Enabled:                   it.enabled,
			AffectNamespaceAllowList:  []string{"*"},
			AffectClusterCreatedAfter: metav1.Time{Time: now},
		}
		pdMixedDeployConfig.Init()
		provider := &cloudv1.Provider{
			Spec: cloudv1.ProviderSpec{
				InstanceTypes: []cloudv1.InstanceType{
					{
						Name: "tikv-instance-type",
						CPU:  resource.MustParse("4"),
					},
				},
			},
		}
		result, _, err := ShouldUseZeroPDNodes(it.pdReplicas, it.pdReplicas, 1, 0, fmt.Sprintf("testing #%d", i), tc,
			provider, nil, mockShootClient, cluster, pdMixedDeployConfig)
		if err != nil {
			t.Fatal("test " + it.name + ", error: " + err.Error())
		}
		if result != it.result {
			t.Fatal("test " + it.name)
		}
	}
}

func TestMixedDeployResourceLimiting(t *testing.T) {
	instances := []cloudv1.InstanceType{
		{
			CPU:    resource.MustParse("2"),
			Memory: resource.MustParse("2Gi"),
		},
		{
			CPU:    resource.MustParse("2"),
			Memory: resource.MustParse("4Gi"),
		},
		{
			CPU:    resource.MustParse("4"),
			Memory: resource.MustParse("8Gi"),
		},
		{
			CPU:    resource.MustParse("8"),
			Memory: resource.MustParse("16Gi"),
		},
		{
			CPU:    resource.MustParse("16"),
			Memory: resource.MustParse("32Gi"),
		},
	}

	tikvRequests := []cloudv1.InstanceType{
		{
			CPU:    resource.MustParse("800m"),
			Memory: resource.MustParse("773094112"), // 773Mi
		},
		{
			CPU:    resource.MustParse("800m"),
			Memory: resource.MustParse("1546188226"), // 1.5Gi
		},
		{
			CPU:    resource.MustParse("2600m"),
			Memory: resource.MustParse("3092376453"), // 3Gi
		},
		{
			CPU:    resource.MustParse("6200m"),
			Memory: resource.MustParse("6184752906"), // 6.1Gi
		},
		{
			CPU:    resource.MustParse("13400m"),
			Memory: resource.MustParse("12369505812"), // 18Gi
		},
	}

	pdRequests := []cloudv1.InstanceType{
		{
			CPU:    resource.MustParse("200m"),
			Memory: resource.MustParse("128849018"), // 128Mi
		},
		{
			CPU:    resource.MustParse("200m"),
			Memory: resource.MustParse("257698037"), // 257Mi
		},
		{
			CPU:    resource.MustParse("400m"),
			Memory: resource.MustParse("515396075"), // 515Mi
		},
		{
			CPU:    resource.MustParse("800m"),
			Memory: resource.MustParse("1030792151"), // 1Gi
		},
		{
			CPU:    resource.MustParse("1600m"),
			Memory: resource.MustParse("2061584302"), // 2Gi
		},
	}

	pdLimits := []cloudv1.InstanceType{
		{
			CPU:    resource.MustParse("400m"),
			Memory: resource.MustParse("214748364"), // 214Mi
		},
		{
			CPU:    resource.MustParse("400m"),
			Memory: resource.MustParse("429496729"), // 429Mi
		},
		{
			CPU:    resource.MustParse("800m"),
			Memory: resource.MustParse("858993459"), // 858Mi
		},
		{
			CPU:    resource.MustParse("1600m"),
			Memory: resource.MustParse("1717986918"), // 1.7Gi
		},
		{
			CPU:    resource.MustParse("3200m"),
			Memory: resource.MustParse("3435973836"), // 3.4Gi
		},
	}

	cmpQuantity := func(resName string, caseIdx int, a resource.Quantity, b resource.Quantity) {
		if !a.Equal(b) {
			t.Fatalf("case #%d, %s: %v != %v", caseIdx, resName, a.String(), b.String())
		}
	}

	for i, instance := range instances {
		tc := &v1alpha1.TidbCluster{
			Spec: v1alpha1.TidbClusterSpec{
				TiKV: &v1alpha1.TiKVSpec{
					ResourceRequirements: corev1.ResourceRequirements{
						Requests: corev1.ResourceList{},
						Limits:   corev1.ResourceList{},
					},
				},
				PD: &v1alpha1.PDSpec{
					ResourceRequirements: corev1.ResourceRequirements{
						Requests: corev1.ResourceList{},
						Limits:   corev1.ResourceList{},
					},
				},
			},
		}
		tc.Spec.TiKV.Requests[corev1.ResourceCPU] = *resource.NewQuantity(instance.CPU.Value()-1, resource.DecimalSI)
		tc.Spec.TiKV.Requests[corev1.ResourceMemory] = *resource.NewQuantity(int64(float64(instance.Memory.Value())*0.6), resource.BinarySI)
		tc.Spec.PD.Requests[corev1.ResourceCPU] = *resource.NewQuantity(instance.CPU.Value()-1, resource.DecimalSI)
		tc.Spec.PD.Requests[corev1.ResourceMemory] = *resource.NewQuantity(int64(float64(instance.Memory.Value())*0.6), resource.BinarySI)

		computeResourcesInMixedDeploy(tc, instance, true, true)

		cmpQuantity("PD CPU requests", i, tc.Spec.PD.Requests[corev1.ResourceCPU], pdRequests[i].CPU)
		cmpQuantity("PD CPU limits", i, tc.Spec.PD.Limits[corev1.ResourceCPU], pdLimits[i].CPU)
		cmpQuantity("PD memory request", i, tc.Spec.PD.Requests[corev1.ResourceMemory], pdRequests[i].Memory)
		cmpQuantity("PD memory limits", i, tc.Spec.PD.Limits[corev1.ResourceMemory], pdLimits[i].Memory)

		cmpQuantity("TiKV CPU requests", i, tc.Spec.TiKV.Requests[corev1.ResourceCPU], tikvRequests[i].CPU)
		cmpQuantity("TiKV memory request", i, tc.Spec.TiKV.Requests[corev1.ResourceMemory], tikvRequests[i].Memory)
	}
}