package k8s

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/types"
	k8Yaml "k8s.io/apimachinery/pkg/util/yaml"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-api/pkg/pingcap/label"
)

func TestGetComponentTopology(t *testing.T) {
	cases := map[string]struct {
		pods   []*corev1.Pod
		nodes  []*corev1.Node
		expect *ComponentTopology
	}{
		"Normal": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-0",
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-1",
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-2",
				},
			}},
			nodes: []*corev1.Node{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						ZoneLabel: "us-west-2a",
					},
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						ZoneLabel: "us-west-2b",
					},
				},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						ZoneLabel: "us-west-2c",
					},
				},
			}},
			expect: &ComponentTopology{
				Replicas: 3,
				Scheduled: map[string]PodNames{
					"us-west-2a": []string{"tikv-0"},
					"us-west-2b": []string{"tikv-1"},
					"us-west-2c": []string{"tikv-2"},
				},
			},
		},
		"PendingPod": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-1",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}, {
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-2",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{},
			}},
			expect: &ComponentTopology{
				Replicas:  3,
				Scheduled: map[string]PodNames{},
				UnScheduled: []string{
					"tikv-0",
					"tikv-1",
					"tikv-2",
				},
			},
		},
		"NoPod": {
			pods: nil,
			expect: &ComponentTopology{
				Replicas:  0,
				Scheduled: map[string]PodNames{},
			},
		},
		"LegacyNodeLabel": {
			pods: []*corev1.Pod{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						label.ComponentLabelKey: "tikv",
						label.InstanceLabelKey:  "test",
					},
				},
				Spec: corev1.PodSpec{
					NodeName: "tikv-0",
				},
			}},
			nodes: []*corev1.Node{{
				ObjectMeta: metav1.ObjectMeta{
					Name: "tikv-0",
					Labels: map[string]string{
						LegacyZoneLabel: "us-west-2a",
					},
				},
			}},
			expect: &ComponentTopology{
				Replicas: 1,
				Scheduled: map[string]PodNames{
					"us-west-2a": []string{"tikv-0"},
				},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			ctx := context.TODO()
			cb := fake.NewClientBuilder()
			for i := range tc.pods {
				cb.WithObjects(tc.pods[i])
			}
			for i := range tc.nodes {
				cb.WithObjects(tc.nodes[i])
			}
			client := cb.Build()
			result, err := GetComponentTopology(ctx, client, types.NamespacedName{Namespace: "", Name: "test"}, cloudv1.ComponentTiKV)
			g.Expect(err).To(Succeed())
			g.Expect(result).To(Equal(tc.expect))
		})
	}
}

func TestReadZoneOfPV(t *testing.T) {
	cases := map[string]struct {
		pvYaml     string
		expectZone string
	}{
		"gcpDisk": {
			pvYaml: `
apiVersion: v1
kind: PersistentVolume
metadata:
  annotations:
    kubernetes.io/createdby: gce-pd-dynamic-provisioner
    pv.kubernetes.io/bound-by-controller: "yes"
    pv.kubernetes.io/provisioned-by: kubernetes.io/gce-pd
    tidb.pingcap.com/pod-name: db-tidb-1
  labels:
    app.kubernetes.io/component: tidb
    app.kubernetes.io/instance: db
    app.kubernetes.io/managed-by: tidb-operator
    app.kubernetes.io/name: tidb-cluster
    app.kubernetes.io/namespace: tidb1379661944637134178
    failure-domain.beta.kubernetes.io/region: us-east4
    failure-domain.beta.kubernetes.io/zone: us-east4-a
    tidb.pingcap.com/cluster-id: "7173695469332910153"
  name: pvc-0d6440df-acb3-4102-9a3d-7c7a12981680
spec:
  nodeAffinity:
    required:
      nodeSelectorTerms:
      - matchExpressions:
        - key: failure-domain.beta.kubernetes.io/zone
          operator: In
          values:
          - us-east4-a
        - key: failure-domain.beta.kubernetes.io/region
          operator: In
          values:
          - us-east4`,
			expectZone: "us-east4-a",
		},
		"ebs": {
			pvYaml: `apiVersion: v1
kind: PersistentVolume
metadata:
  annotations:
    "pv.kubernetes.io/provisioned-by": "ebs.csi.aws.com"
  name: pv-shoot--prod--fe5c8299-f9b947ba-527d-47df-9659-79286f9da663
spec:
  nodeAffinity:
    required:
      nodeSelectorTerms:
      - matchExpressions:
        - key: topology.ebs.csi.aws.com/zone
          operator: In
          values:
          - us-east-1a`,
			expectZone: "us-east-1a",
		},
	}

	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			pv := &corev1.PersistentVolume{}
			g.Expect(k8Yaml.Unmarshal([]byte(tc.pvYaml), pv)).To(Succeed())

			result := readZoneOfPV(pv)
			g.Expect(result).To(Equal(tc.expectZone))
		})
	}
}