package cluster

import (
	"testing"

	. "github.com/onsi/gomega"

	"github.com/tidbcloud/infra-api/pkg/k8s"
)

func TestScaleInOnePod(t *testing.T) {
	cases := map[string]struct {
		zones   []string
		spreads k8s.PodsByZone
		err     bool
		pod     *statefulPod
	}{
		"ScaleInMaxZone": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-1", "db-tikv-2"},
				"b": []string{"db-tikv-3"},
				"c": []string{"db-tikv-4"},
			},
			pod: &statefulPod{
				name:    "db-tikv-2",
				ordinal: 2,
			},
		},
		"ScaleInLastZoneOnTie": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-2"},
				"b": []string{"db-tikv-1"},
				"c": []string{"db-tikv-0"},
			},
			pod: &statefulPod{
				name:    "db-tikv-0",
				ordinal: 0,
			},
		},
		"ScaleInMaxOrdinalInMaxZone": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-6", "db-tikv-0", "db-tikv-3"},
				"b": []string{"db-tikv-1", "db-tikv-4"},
				"c": []string{"db-tikv-2", "db-tikv-5"},
			},
			pod: &statefulPod{
				name:    "db-tikv-6",
				ordinal: 6,
			},
		},
		"NoPodInOneZone": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-6", "db-tikv-0", "db-tikv-3"},
				"c": []string{"db-tikv-2", "db-tikv-5"},
			},
			pod: &statefulPod{
				name:    "db-tikv-6",
				ordinal: 6,
			},
		},
		"NoPodInOtherZones": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-6", "db-tikv-0", "db-tikv-3"},
			},
			pod: &statefulPod{
				name:    "db-tikv-6",
				ordinal: 6,
			},
		},
		"MoreZones": {
			zones: []string{"a", "b", "c", "d", "e"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-6", "db-tikv-0", "db-tikv-3"},
				"b": []string{"db-tikv-1", "db-tikv-4"},
				"c": []string{"db-tikv-2", "db-tikv-5"},
				"e": []string{"db-tikv-7", "db-tikv-9"},
				"f": []string{"db-tikv-8", "db-tikv-10"},
			},
			pod: &statefulPod{
				name:    "db-tikv-6",
				ordinal: 6,
			},
		},
		"NoPodToScale": {
			zones:   []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{},
			err:     true,
		},
		"IllegalPod": {
			zones: []string{"a", "b", "c"},
			spreads: map[string]k8s.PodNames{
				"a": []string{"db-tikv-zzzz"},
				"b": []string{"db-tikv-aaaa"},
				"c": []string{"db-tikv-cccc"},
			},
			err: true,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			pod, err := scaleInOnePod(tc.zones, tc.spreads)
			if tc.err {
				g.Expect(err).NotTo(Succeed())
			} else {
				g.Expect(err).To(Succeed())
				g.Expect(pod).To(Equal(tc.pod))
			}
		})
	}
}

func TestParseStatefulSetPod(t *testing.T) {
	cases := map[string]struct {
		name string
		err  bool
		pod  *statefulPod
	}{
		"TiKV": {
			name: "db-tikv-1",
			pod: &statefulPod{
				name:    "db-tikv-1",
				ordinal: 1,
			},
		},
		"Complex": {
			name: "a-b-c-d-e-f-666",
			pod: &statefulPod{
				name:    "a-b-c-d-e-f-666",
				ordinal: 666,
			},
		},
		"DeploymentPod": {
			name: "nginx-eeeeee-eeeeee",
			err:  true,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			pod, err := parseStatefulPod(tc.name)
			if tc.err {
				g.Expect(err).NotTo(Succeed())
			} else {
				g.Expect(err).To(Succeed())
				g.Expect(pod).To(Equal(tc.pod))
			}
		})
	}
}