package tidbpdb

import (
	"context"
	"fmt"
	"testing"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"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"
	corev1 "k8s.io/api/core/v1"
	policyv1 "k8s.io/api/policy/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/utils/ptr"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func TestTidbPDB_do(t *testing.T) {
	tests := []struct {
		name             string
		freeCacheSize    int
		tidbPoolMinSize  int
		readyPods        int
		expectedMinAvail *int
	}{
		{
			name:             "drain free cache",
			freeCacheSize:    6,
			tidbPoolMinSize:  3,
			readyPods:        8,
			expectedMinAvail: ptr.To(3),
		},
		{
			name:             "no free cache",
			freeCacheSize:    0,
			tidbPoolMinSize:  2,
			readyPods:        3,
			expectedMinAvail: ptr.To(3),
		},
		{
			name:             "tidbPoolMinSize < 2",
			freeCacheSize:    2,
			tidbPoolMinSize:  1,
			readyPods:        2,
			expectedMinAvail: ptr.To(0),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			objs := make([]runtime.Object, tt.readyPods)
			for i := 0; i < tt.readyPods; i++ {
				pod := &corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      fmt.Sprintf("tidb-server-%d", i),
						Namespace: "tidb-pool",
						Labels:    tidbMatchLabels,
					},
					Status: corev1.PodStatus{Conditions: []corev1.PodCondition{
						{
							Type:   corev1.PodReady,
							Status: corev1.ConditionTrue,
						},
					}},
				}
				objs[i] = pod
			}
			objs = append(objs, &corev1.ConfigMap{
				ObjectMeta: metav1.ObjectMeta{
					Name:      types.ServerlessPoolConfig,
					Namespace: "tidb-pool",
				},
				Data: map[string]string{types.MinPoolReplicasKey: fmt.Sprint(tt.tidbPoolMinSize)},
			})
			tidbPDB := &TidbPDB{
				Logger:            logr.New(log.NullLogSink{}),
				client:            fake.NewClientBuilder().WithScheme(scheme.Scheme).WithRuntimeObjects(objs...).Build(),
				tManager:          &clientfake.FakeTidbManager{FreeCacheSizeR: tt.freeCacheSize},
				tidbPoolNamespace: "tidb-pool",
			}

			tidbPDB.do()

			pdb := &policyv1.PodDisruptionBudget{}
			err := tidbPDB.client.Get(context.Background(), client.ObjectKey{Name: tidbPDBName, Namespace: "tidb-pool"}, pdb)
			if tt.expectedMinAvail == nil {
				g.Expect(err).To(HaveOccurred())
				g.Expect(apierrors.IsNotFound(err)).To(BeTrue())
				return
			}
			g.Expect(err).NotTo(HaveOccurred())
			g.Expect(pdb.Spec.MinAvailable.IntValue()).To(Equal(*tt.expectedMinAvail))
		})
	}
}