package tidbmanager

import (
	"context"
	"testing"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/pingcap/errors"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/clientfake"
	"go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.etcd.io/etcd/api/v3/mvccpb"
	clientv3 "go.etcd.io/etcd/client/v3"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func Test_tidbManager_FreeTidb(t1 *testing.T) {
	tidbPoolNamespace := "tidb-pool"
	type args struct {
		podName string
		podIP   string
	}
	tests := []struct {
		name                   string
		args                   args
		pods                   []runtime.Object
		unpatchCriticalTidbPod bool
		txnResp                clientv3.TxnResponse
		txnErr                 error
		assignCache            map[string]tidbInfo
		wantErr                error
	}{
		{
			name:    "FreeTidb, free pod1,ip1 etcd resp Succeeded",
			args:    args{podName: "pod1", podIP: "ip1"},
			txnResp: clientv3.TxnResponse{Succeeded: true},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}},
			},
		},
		{
			name: "FreeTidb, free pod1,ip1 etcd resp false",
			args: args{podName: "pod1", podIP: "ip1"},
			txnResp: clientv3.TxnResponse{
				Succeeded: false,
				Responses: []*etcdserverpb.ResponseOp{{
					Response: &etcdserverpb.ResponseOp_ResponseRange{ResponseRange: &etcdserverpb.RangeResponse{
						Kvs: []*mvccpb.KeyValue{{Value: []byte("pod2,ip2")}},
					}},
				}},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}},
			},
		},
		{
			name: "FreeTidb, free pod1,ip1 etcd resp no kvs",
			args: args{podName: "pod1", podIP: "ip1"},
			txnResp: clientv3.TxnResponse{
				Succeeded: false,
				Responses: []*etcdserverpb.ResponseOp{{}},
			},
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}},
			},
		},
		{
			name:   "FreeTidb, free pod1,ip1 txn error",
			args:   args{podName: "pod1", podIP: "ip1"},
			txnErr: errors.New("txn error"),
			assignCache: map[string]tidbInfo{
				"token1": {endpoints: []string{"pod1,ip1", "pod2,ip2"}},
			},
			wantErr: errors.New("txn error"),
		},
	}
	for _, tt := range tests {
		t1.Run(tt.name, func(t1 *testing.T) {
			g := NewGomegaWithT(t1)

			pod := &corev1.Pod{
				ObjectMeta: metav1.ObjectMeta{Name: tt.args.podName, Namespace: tidbPoolNamespace},
				Status:     corev1.PodStatus{PodIP: tt.args.podIP},
			}
			if tt.unpatchCriticalTidbPod {
				pod.Annotations = map[string]string{types.CASafeToEvict: "false"}
			}
			b := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(pod)
			serverlessPoolConfig := &corev1.ConfigMap{ObjectMeta: metav1.ObjectMeta{
				Name:      types.ServerlessPoolConfig,
				Namespace: tidbPoolNamespace,
			}, Data: map[string]string{}}
			if tt.unpatchCriticalTidbPod {
				serverlessPoolConfig.Data[types.TidbCriticalKey] = "true"
			}
			b.WithRuntimeObjects(serverlessPoolConfig)

			t := &tidbManager{
				Logger:            logr.New(log.NullLogSink{}),
				client:            b.Build(),
				ctx:               context.Background(),
				safeToEvict:       types.CASafeToEvict,
				tidbPoolNamespace: tidbPoolNamespace,
				etcdKV:            clientfake.NewEtcdKV(clientfake.BuildArg{TxnResponse: &tt.txnResp, TxnErr: tt.txnErr}),
				etcdLease:         clientfake.NewEtcdLease(),
				freeCache:         &freeCache{},
				assignCache:       &assignCache{},
			}
			t.assignCache.init(tt.assignCache)

			err := t.FreeTidb(context.Background(), tt.args.podName, tt.args.podIP, "")
			if tt.wantErr == nil {
				g.Expect(err).To(Succeed())
				if tt.unpatchCriticalTidbPod {
					pod := &corev1.Pod{}
					g.Expect(t.client.Get(context.Background(), client.ObjectKey{
						Namespace: tidbPoolNamespace,
						Name:      tt.args.podName,
					}, pod)).To(Succeed())
					g.Expect(pod.Annotations).NotTo(HaveKey(t.safeToEvict))
				}
			} else {
				g.Expect(err).To(MatchError(tt.wantErr.Error()))
			}
		})
	}
}