package zoneredundantevicter

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/tidbmanager"
	"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"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
)

func TestReconciler_evictZoneRedundantTidbPods(t *testing.T) {
	tests := []struct {
		name         string
		tidbManagers map[string]tidbmanager.TidbManager
		evictNodes   []*corev1.Node
		objs         []client.Object
		infraObjs    []client.Object
		evictErr     error
		wantRetry    bool
	}{
		{
			name: "evict tidb pods, no retry",
			tidbManagers: map[string]tidbmanager.TidbManager{
				"tidb-pool": &clientfake.FakeTidbManager{
					GetTidbFromFreeCacheR: map[string][]string{
						"node1": {"tidb-server-fc-node1-pod1"},
						"node2": {"tidb-server-fc-node2-pod1"},
					},
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"token1": {Endpoints: []string{
							"tidb-server-assigned-node1-pod1,ip1",
							"tidb-server-assigned-node3-not-evict-pod2,ip2",
						}},
						"token2": {Endpoints: []string{
							"tidb-server-assigned-node2-pod2,ip2",
						}},
						"token-not-zr": {Endpoints: []string{
							"tidb-server-not-zr-node1-not-evict,ip3",
						}},
					},
				},
			},
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-server-not-zr-node1-not-evict",
						Namespace: "tidb-pool",
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-server-other-node3-not-evict",
						Namespace: "tidb-pool",
					},
					Spec: corev1.PodSpec{
						NodeName: "node3",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-server-assigned-node3-not-evict-pod2",
						Namespace: "tidb-pool",
					},
					Spec: corev1.PodSpec{
						NodeName: "node3",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "tidb-server-fc-node1-pod1",
						Namespace:   "tidb-pool",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "tidb-server-fc-node2-pod1",
						Namespace:   "tidb-pool",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "tidb-server-assigned-node1-pod1",
						Namespace:   "tidb-pool",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "tidb-server-assigned-node2-pod2",
						Namespace:   "tidb-pool",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
			},
			infraObjs: []client.Object{
				&cloudv1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: "cluster-not-zr", Namespace: "cluster-project1"},
					Status: cloudv1.ClusterStatus{
						FromProvider: runtime.RawExtension{
							Raw: marshal(&cloudv1.AwsFreeTierClusterStatus{
								TidbPool:      "tidb-pool",
								KeyspaceToken: "token-not-zr",
							}),
						},
					},
				},
				&cloudv1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: "cluster-token1", Namespace: "cluster-project2"},
					Status: cloudv1.ClusterStatus{
						FromProvider: runtime.RawExtension{
							Raw: marshal(&cloudv1.AwsFreeTierClusterStatus{
								TidbPool:             "tidb-pool",
								KeyspaceToken:        "token1",
								ZoneRedundantEnabled: true,
							}),
						},
					},
				},
				&cloudv1.Branch{
					ObjectMeta: metav1.ObjectMeta{Name: "branch-token2", Namespace: "cluster-project3"},
					Status: cloudv1.BranchStatus{
						FromProvider: runtime.RawExtension{
							Raw: marshal(&cloudv1.AwsFreeTierClusterStatus{
								TidbPool:             "tidb-pool",
								KeyspaceToken:        "token2",
								ZoneRedundantEnabled: true,
							}),
						},
					},
				},
			},
		},
		{
			name: "evict tidb pods, retry",
			tidbManagers: map[string]tidbmanager.TidbManager{
				"tidb-pool": &clientfake.FakeTidbManager{
					GetAssignedTidbByTokenR: map[string]*ep.Endpoints{
						"token1": {Endpoints: []string{
							"tidb-server-assigned-node1-not-evict-pod,ip1",
							"tidb-server-assigned-node2-pod,ip2",
						}},
					},
				},
			},
			evictNodes: []*corev1.Node{
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node1"},
					Status: corev1.NodeStatus{Conditions: []corev1.NodeCondition{{
						Type:   corev1.NodeReady,
						Status: corev1.ConditionTrue,
					}}},
				},
				{
					ObjectMeta: metav1.ObjectMeta{Name: "node2"},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{Key: corev1.TaintNodeUnreachable, Effect: corev1.TaintEffectNoExecute}},
					},
				},
			},
			objs: []client.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "tidb-server-assigned-node1-not-evict-pod",
						Namespace: "tidb-pool",
					},
					Spec: corev1.PodSpec{
						NodeName: "node1",
					},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:        "tidb-server-assigned-node2-pod",
						Namespace:   "tidb-pool",
						Annotations: map[string]string{annoObjectShouldBeDeletedKey: "1"},
					},
					Spec: corev1.PodSpec{
						NodeName: "node2",
					},
				},
			},
			infraObjs: []client.Object{
				&cloudv1.Cluster{
					ObjectMeta: metav1.ObjectMeta{Name: "cluster-token1", Namespace: "cluster-project1"},
					Status: cloudv1.ClusterStatus{
						FromProvider: runtime.RawExtension{
							Raw: marshal(&cloudv1.AwsFreeTierClusterStatus{
								TidbPool:             "tidb-pool",
								KeyspaceToken:        "token1",
								ZoneRedundantEnabled: true,
							}),
						},
					},
				},
			},
			evictErr:  apierrors.NewTooManyRequestsError("inject to ut"),
			wantRetry: true,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			allObjs := make([]client.Object, 0, len(tt.objs)+len(tt.evictNodes))
			allObjs = append(allObjs, tt.objs...)
			for _, n := range tt.evictNodes {
				allObjs = append(allObjs, n)
			}
			r := newFakeReconciler(allObjs, tt.infraObjs, tt.tidbManagers, tt.evictErr)
			gotRetry, err := r.evictZoneRedundantTidbPods(context.Background(), tt.evictNodes)
			g.Expect(err).To(Succeed())
			g.Expect(gotRetry).To(Equal(tt.wantRetry))
			checkNodesAndObjects(g, r.Client, tt.evictNodes, tt.objs)
		})
	}
}