package branch

import (
	"encoding/json"
	"strconv"
	"strings"
	"testing"
	"time"

	. "github.com/onsi/gomega"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	ep "github.com/tidbcloud/aws-shared-provider/pkg/utils/endpoint"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
)

func marshal(v any) []byte {
	b, _ := json.Marshal(v)
	return b
}

func TestGenPublicEndpointStatus(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	cases := map[string]struct {
		branch         *cloudv1.Branch
		observed       []runtime.Object
		expectEndpoint *cloudv1.SqlEndpointStatus
	}{
		"CreatePublicEndpoint": {
			branch: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "Test-CreatePublicEndpoint",
					Namespace: "test-ns",
				},
				Spec: cloudv1.BranchSpec{
					SqlEndpoints: cloudv1.BranchSqlEndpointsSpec{
						Public: cloudv1.BranchPublicSpec{
							Service: cloudv1.Service{
								Enabled: ptr.To(true),
							},
						},
					},
				},
			},
			observed: []runtime.Object{
				&corev1.Service{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-gateway-new",
						Namespace: "tidb-admin",
						Annotations: map[string]string{
							types.AnnoGatewayDnsName: "gateway01.us-east-1.dev.shared.aws.tidbcloud.com",
						},
					},
				},
			},
			expectEndpoint: &cloudv1.SqlEndpointStatus{
				Host: "gateway01.us-east-1.dev.shared.aws.tidbcloud.com",
				Port: defaultTiDBPort,
			},
		},
	}

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

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

			r := newFakeReconcileContext(c.branch, newFakeReconciler(c.observed, []runtime.Object{c.branch}, httpServer), nil)
			publicEndpoint, err := r.genPublicEndpointStatus()
			g.Expect(err).To(Succeed())
			g.Expect(publicEndpoint).To(Equal(c.expectEndpoint))
		})
	}
}

func TestGenPrivatelinkStatus(t *testing.T) {
	t.Setenv("NAMESPACE", "tidb-admin")

	cases := map[string]struct {
		cluster           *cloudv1.Branch
		observed          []runtime.Object
		expectPrivatelink *cloudv1.PrivateLinkServiceStatus
	}{
		"CreatePrivatelink": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "Test-CreatePrivatelink",
					Namespace: "test-ns",
				},
				Spec: cloudv1.BranchSpec{
					SqlEndpoints: cloudv1.BranchSqlEndpointsSpec{
						PrivateLinkService: &cloudv1.PrivateLinkService{
							AWS: &cloudv1.AWSPrivateLinkService{AcceptedEndpoints: []string{"*"}},
						},
					},
				},
			},
			observed: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "privatelink-service-config",
						Namespace: "tidb-admin",
					},
					Data: map[string]string{
						"private-dns":        "gateway01-privatelink.us-east-1.dev.shared.aws.tidbcloud.com",
						"service-name":       "com.amazonaws.vpce.us-east-1.vpce-svc-0699cfb9b1db28a22",
						"availability-zones": "us-east-1a,us-west-2b",
					},
				},
			},
			expectPrivatelink: &cloudv1.PrivateLinkServiceStatus{
				Phase: cloudv1.PrivateLinkServicePhaseActive,
				AWS: &cloudv1.AWSPrivateLinkServiceStatus{
					AvailabilityZones: []string{"us-east-1a", "us-west-2b"},
					ServiceName:       "com.amazonaws.vpce.us-east-1.vpce-svc-0699cfb9b1db28a22",
					PrivateDNSName:    "gateway01-privatelink.us-east-1.dev.shared.aws.tidbcloud.com",
					Port:              ptr.To[int32](defaultTiDBPort),
				},
			},
		},
		"CreatePrivatelinkZR": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "Test-CreatePrivatelink-ZR",
					Namespace: "test-ns",
				},
				Spec: cloudv1.BranchSpec{
					SqlEndpoints: cloudv1.BranchSqlEndpointsSpec{
						PrivateLinkService: &cloudv1.PrivateLinkService{
							AWS: &cloudv1.AWSPrivateLinkService{AcceptedEndpoints: []string{"*"}},
						},
					},
					ForProvider: &runtime.RawExtension{
						Raw: marshal(&types.InfraServerlessSpec{EnableZoneRedundant: true}),
					},
				},
			},
			observed: []runtime.Object{
				&corev1.ConfigMap{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "privatelink-zr-service-config",
						Namespace: "tidb-admin",
					},
					Data: map[string]string{
						"private-dns":        "gateway01-privatelink-zr.us-east-1.dev.shared.aws.tidbcloud.com",
						"service-name":       "com.amazonaws.vpce.us-east-1.vpce-svc-9999cfb9c1db28a88",
						"availability-zones": "us-east-1c,us-west-2b",
					},
				},
			},
			expectPrivatelink: &cloudv1.PrivateLinkServiceStatus{
				Phase: cloudv1.PrivateLinkServicePhaseActive,
				AWS: &cloudv1.AWSPrivateLinkServiceStatus{
					AvailabilityZones: []string{"us-east-1c", "us-west-2b"},
					ServiceName:       "com.amazonaws.vpce.us-east-1.vpce-svc-9999cfb9c1db28a88",
					PrivateDNSName:    "gateway01-privatelink-zr.us-east-1.dev.shared.aws.tidbcloud.com",
					Port:              ptr.To[int32](defaultTiDBPort),
				},
			},
		},
	}

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

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

			r := newFakeReconcileContext(c.cluster, newFakeReconciler(c.observed, []runtime.Object{c.cluster}, httpServer), nil)
			specForProvider, err := types.InfraServerlessSpecForProvider(c.cluster)
			g.Expect(err).To(Succeed())
			privatelink, err := r.genPrivatelinkStatus(specForProvider)
			g.Expect(err).To(Succeed())
			g.Expect(privatelink).To(Equal(c.expectPrivatelink))
		})
	}
}

func TestReconcileTidbPool(t *testing.T) {
	cases := map[string]struct {
		cluster               *cloudv1.Branch
		tidbManagers          map[string]tidbmanager.TidbManager
		etcdTokenPrefixMap    map[string]string
		expectAnno            bool
		expectOldPoolReplicas int
		expectErr             error
	}{
		"PoolUnchanged": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:        "test",
					Namespace:   "ns",
					Annotations: map[string]string{types.AnnoTidbPoolSelector: defaultTidbPool},
				},
				Spec: cloudv1.BranchSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{
							Replicas:     1,
							InstanceType: "tidb.freetier",
						},
					},
				},
				Status: cloudv1.BranchStatus{
					Phase: cloudv1.BranchPhaseRunning,
					DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
					FromProvider: runtime.RawExtension{
						Object: &cloudv1.AwsFreeTierClusterStatus{TidbPool: defaultTidbPool},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
			},
		},
		"PoolStartChanging": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:        "test",
					Namespace:   "ns",
					Annotations: map[string]string{types.AnnoTidbPoolSelector: "new-pool"},
				},
				Spec: cloudv1.BranchSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{
							Replicas:     2,
							InstanceType: "tidb.freetier",
						},
					},
				},
				Status: cloudv1.BranchStatus{
					Phase: cloudv1.BranchPhaseRunning,
					DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 2}},
					FromProvider: runtime.RawExtension{
						Object: &cloudv1.AwsFreeTierClusterStatus{KeyspaceToken: "keyspace-token", TidbPool: defaultTidbPool},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: 2,
			expectErr:             &utils.RequeueAfterError{},
		},
		"PoolChanging_WaitNewTidbPoolRunning": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
						types.AnnoTidbPoolSelector:    "new-pool",
					},
				},
				Spec: cloudv1.BranchSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{
							Replicas:     2,
							InstanceType: "tidb.freetier",
						},
					},
				},
				Status: cloudv1.BranchStatus{
					Phase: cloudv1.BranchPhaseRunning,
					DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 2}},
					FromProvider: runtime.RawExtension{
						Object: &cloudv1.AwsFreeTierClusterStatus{KeyspaceToken: "keyspace-token", TidbPool: "new-pool"},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: 2,
			expectErr: &utils.RequeueAfterError{
				Cause:        errors.New("waiting new tidb pool pods to running"),
				RequeueAfter: waitNewTidbPoolRunning,
			},
		},
		"PoolChanging_ReduceOne": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Add(-retryTidbPoolChanging).Format(time.RFC3339),
						types.AnnoTidbPoolSelector:    "new-pool",
					},
				},
				Spec: cloudv1.BranchSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{
							Replicas:     3,
							InstanceType: "tidb.freetier",
						},
					},
				},
				Status: cloudv1.BranchStatus{
					Phase: cloudv1.BranchPhaseRunning,
					DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 3}},
					FromProvider: runtime.RawExtension{
						Object: &cloudv1.AwsFreeTierClusterStatus{KeyspaceToken: "keyspace-token", TidbPool: "new-pool"},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1", "tidb-server-pod4,ip4"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod2,ip2"}},
					},
				},
				"new-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod3,ip3"}},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
				"new-pool":      "/new-pool/",
			},
			expectAnno:            true,
			expectOldPoolReplicas: 2,
			expectErr:             &utils.RequeueAfterError{RequeueAfter: retryTidbPoolChanging},
		},
		"PoolEndChanging": {
			cluster: &cloudv1.Branch{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "ns",
					Annotations: map[string]string{
						types.AnnoKeyTidbPoolChanging: time.Now().Format(time.RFC3339),
						types.AnnoTidbPoolSelector:    defaultTidbPool,
					},
				},
				Spec: cloudv1.BranchSpec{
					Cluster: cloudv1.ClusterTopology{
						TiDB: cloudv1.TiDB{
							Replicas:     1,
							InstanceType: "tidb.freetier",
						},
					},
				},
				Status: cloudv1.BranchStatus{
					Phase: cloudv1.BranchPhaseRunning,
					DB:    &cloudv1.DBStatus{TiDB: &cloudv1.TiDBStatus{UpdatedReplicas: 1}},
					FromProvider: runtime.RawExtension{
						Object: &cloudv1.AwsFreeTierClusterStatus{KeyspaceToken: "keyspace-token", TidbPool: defaultTidbPool},
					},
				},
			},
			tidbManagers: map[string]tidbmanager.TidbManager{
				defaultTidbPool: &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {Endpoints: []string{"tidb-server-pod1,ip1"}},
					},
				},
				"other-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"keyspace-token": {},
					},
				},
			},
			etcdTokenPrefixMap: map[string]string{
				defaultTidbPool: "/default-pool/",
				"other-pool":    "/other-pool/",
			},
		},
	}

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

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

			tidbPods := make([]runtime.Object, 0)
			etcdDB := make(map[string]string)
			for pool, tManager := range c.tidbManagers {
				for keyspaceToken, eps := range tManager.(*clientfake.FakeTidbManager).GetAssignedTidbByTokenR {
					etcdDB[c.etcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyReplicas] = strconv.Itoa(len(eps.Endpoints))
					etcdDB[c.etcdTokenPrefixMap[pool]+keyspaceToken+"/"+utils.KeyEndpoints] = strings.Join(eps.Endpoints, ";")
					for _, endpoint := range eps.Endpoints {
						podName, podIP := ep.ParsePodFromEndpoint(endpoint)
						tidbPod := &corev1.Pod{
							ObjectMeta: metav1.ObjectMeta{
								Name:        podName,
								Namespace:   pool,
								Annotations: make(map[string]string),
							},
							Status: corev1.PodStatus{Phase: corev1.PodRunning, PodIP: podIP},
						}
						// intentionally setting up AnnoTidbStatusRunSlow in the new pool
						// if pool == c.cluster.Annotations[types.AnnoTidbPoolSelector] {
						// 	tidbPod.Annotations[types.AnnoTidbStatusRunSlow] = "true"
						// }
						tidbPods = append(tidbPods, tidbPod)
					}
				}
			}

			branchStatus, err := c.cluster.AwsFreeTierClusterStatusFromBranch()
			g.Expect(err).To(Succeed())

			nr := newFakeReconciler(tidbPods, []runtime.Object{c.cluster}, httpServer)
			nr.TidbManagers = c.tidbManagers
			nr.EtcdKV = clientfake.NewEtcdKV(clientfake.BuildArg{DB: etcdDB})
			nr.EtcdTokenPrefixMap = c.etcdTokenPrefixMap
			nr.TidbPools = []string{defaultTidbPool}
			for pool := range c.etcdTokenPrefixMap {
				if pool != defaultTidbPool {
					nr.TidbPools = append(nr.TidbPools, pool)
				}
			}
			r := newFakeReconcileContext(c.cluster, nr, nil)
			r.etcdTokenPrefix = c.etcdTokenPrefixMap[c.cluster.Annotations[types.AnnoTidbPoolSelector]]

			err = r.reconcileTidbPool()
			if c.expectErr != nil {
				if reqErr, ok := c.expectErr.(*utils.RequeueAfterError); ok && reqErr.RequeueAfter == 0 {
					g.Expect(err.(*utils.RequeueAfterError).RequeueAfter).To(BeNumerically("<", retryTidbPoolChanging))
				} else {
					g.Expect(err).To(MatchError(c.expectErr.Error()))
				}
			} else {
				g.Expect(err).To(Succeed())
			}
			if c.expectAnno {
				g.Expect(r.branch.Annotations).To(HaveKey(types.AnnoKeyTidbPoolChanging))

				// verify the all old pools replicas
				oldPoolReplicas := 0
				for _, etcdTokenPrefix := range c.etcdTokenPrefixMap {
					if etcdTokenPrefix == r.etcdTokenPrefix {
						continue
					}
					resp, err := r.EtcdKV.Get(r.ctx, etcdTokenPrefix+r.keyspaceToken+"/"+utils.KeyReplicas)
					g.Expect(err).To(Succeed())
					if resp.Count != 0 {
						r, err := strconv.Atoi(string(resp.Kvs[0].Value))
						g.Expect(err).To(Succeed())
						oldPoolReplicas += r
					}
				}
				g.Expect(oldPoolReplicas).To(Equal(c.expectOldPoolReplicas))

				if c.cluster.Annotations[types.AnnoTidbPoolSelector] != branchStatus.TidbPool {
					// verify the new pool replicas
					resp, err := r.EtcdKV.Get(r.ctx, r.etcdTokenPrefix+r.keyspaceToken+"/"+utils.KeyReplicas)
					g.Expect(err).To(Succeed())
					g.Expect(string(resp.Kvs[0].Value)).To(Equal(strconv.Itoa(c.cluster.Spec.Cluster.TiDB.Replicas)))

					// verify the old tidb pools pod annotations, should have AnnoTidbStatusRunSlow
					// for _, objKey := range r.allOldTidbPoolPods() {
					// 	tidbPod := &corev1.Pod{}
					// 	g.Expect(r.Client.Get(r.ctx, objKey, tidbPod)).To(Succeed())
					// 	g.Expect(tidbPod.Annotations).To(HaveKey(types.AnnoTidbStatusRunSlow))
					// }

					// verify the new tidb pool pod annotations, should not have AnnoTidbStatusRunSlow
					// for _, objKey := range r.newTidbPoolPods() {
					// 	tidbPod := &corev1.Pod{}
					// 	g.Expect(r.Client.Get(r.ctx, objKey, tidbPod)).To(Succeed())
					// 	g.Expect(tidbPod.Annotations).ToNot(HaveKey(types.AnnoTidbStatusRunSlow))
					// }
				}
			} else {
				g.Expect(r.branch.Annotations).ToNot(HaveKey(types.AnnoKeyTidbPoolChanging))
				// verify all old tidb pool paths are removed
				for _, etcdTokenPrefix := range c.etcdTokenPrefixMap {
					if etcdTokenPrefix == r.etcdTokenPrefix {
						continue
					}
					resp, err := r.EtcdKV.Get(r.ctx, etcdTokenPrefix+r.keyspaceToken+"/"+utils.KeyReplicas)
					g.Expect(err).To(Succeed())
					g.Expect(resp.Count).To(Equal(int64(0)))
				}
			}
		})
	}
}