package cluster

import (
	"strconv"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	infra "github.com/tidbcloud/tidb-management-service/api/serverless_infra"
	"github.com/tidbcloud/tidb-management-service/api/serverless_infra/meta"
	clientv3 "go.etcd.io/etcd/client/v3"
)

func TestEnsureResourceGroup(t *testing.T) {
	cases := map[string]struct {
		cluster              *infra.Cluster
		expectResourceGroups map[string]infra.ResourceGroup
	}{
		"EnsureResourceGroup": {
			cluster: &infra.Cluster{
				ObjectMeta: meta.ObjectMeta{
					Name:        "Test-EnsureResourceGroup",
					Namespace:   "test-ns",
					TagsTenant:  "1",
					TagsProject: "1",
				},
				Spec: infra.ClusterSpec{TiDB: infra.TiDBSpec{
					Replicas: 1,
					ResourceGroups: map[string]*infra.ResourceGroup{
						"rg1": {Replicas: 1},
						"rg3": {Pool: "other-pool", Replicas: 2},
						// not change the content of ResourceGroup by ensureResourceGroup
						"rg4": {Pool: defaultTidbPool, Replicas: 3},
					},
				}},
				Status: infra.ClusterStatus{
					Phase:         infra.ClusterPhaseRunning,
					TidbPool:      defaultTidbPool,
					TiDB:          infra.TiDBStatus{UpdatedReplicas: 1},
					KeyspaceToken: "keyspace-token",
					ResourceGroups: map[string]infra.ResourceGroup{
						"rg1": {Pool: defaultTidbPool, Replicas: 1},
						"rg2": {Pool: "other-pool", Replicas: 1},
						"rg4": {Pool: "other-pool", Replicas: 2},
					},
				},
			},
			expectResourceGroups: map[string]infra.ResourceGroup{
				"rg1": {Pool: defaultTidbPool, Replicas: 1},
				"rg3": {Pool: "other-pool", Replicas: 2},
				"rg4": {Pool: "other-pool", Replicas: 2},
			},
		},
	}

	httpServer := newFakeHTTP()
	defer httpServer.Close()

	for name, c := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			r := newFakeReconcileContext(c.cluster, nil, httpServer)
			etcdDB := make(map[string]string)
			for rgName, rgStatus := range c.cluster.Status.ResourceGroups {
				etcdDB[r.EtcdTokenPrefixMap[rgStatus.Pool]+r.keyspaceToken+"/"+rgName+"/"+utils.KeyReplicas] = strconv.Itoa(rgStatus.Replicas)
				etcdDB[r.EtcdTokenPrefixMap[rgStatus.Pool]+r.keyspaceToken+"/"+rgName+"/"+utils.KeyEndpoints] = ""
			}
			r.EtcdKV = clientfake.NewEtcdKV(clientfake.BuildArg{DB: etcdDB})
			g.Expect(r.ensureResourceGroup()).To(Succeed())
			g.Expect(c.cluster.Status.ResourceGroups).To(Equal(c.expectResourceGroups))

			resp, err := r.EtcdKV.Get(r.ctx, "/", clientv3.WithPrefix())
			g.Expect(err).To(Succeed())
			g.Expect(resp.Count).To(Equal(int64(len(c.expectResourceGroups) * 2)))

			for rgName, rgExpect := range c.expectResourceGroups {
				etcdReplicasKey := r.EtcdTokenPrefixMap[rgExpect.Pool] + r.keyspaceToken + "/" + rgName + "/" + utils.KeyReplicas
				resp, err := r.EtcdKV.Get(r.ctx, etcdReplicasKey)
				g.Expect(err).To(Succeed())
				replicas, err := strconv.Atoi(string(resp.Kvs[0].Value))
				g.Expect(err).To(Succeed())
				g.Expect(replicas).To(Equal(rgExpect.Replicas))
			}
		})
	}
}