package extratidb

import (
	"context"
	"fmt"
	"math/rand"
	"testing"

	oplabel "github.com/pingcap/tidb-operator/pkg/apis/label"
	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/infra-provider/pkg/consts"
)

func getTestNode(clusterID, component, zone string) *corev1.Node {
	return &corev1.Node{
		ObjectMeta: metav1.ObjectMeta{
			Name: fmt.Sprintf("%s-%s-%s-%d", clusterID, component, zone, rand.Int()),
			Labels: map[string]string{
				"tidbcloud-cluster": clusterID,
				"tidb-component":    component,
				consts.ZoneLabelKey: zone,
			},
		},
	}
}

func TestGetExtraTiDBZone(t *testing.T) {
	ctx := context.TODO()
	clusterID := "test"
	zones := []string{"us-west-1a", "us-west-1b", "us-west-1c"}

	tests := []struct {
		name              string
		nodes             []*corev1.Node
		wantExtraTiDBZone string
		wantErr           bool
	}{
		{
			name: "only one tidb node",
			nodes: []*corev1.Node{
				getTestNode(clusterID, "tidb", "us-west-1a"),
				getTestNode(clusterID, "tikv", "us-west-1a"),
				getTestNode(clusterID, "tikv", "us-west-1b"),
				getTestNode(clusterID, "tikv", "us-west-1c"),
			},
			wantExtraTiDBZone: "us-west-1b",
		},
		{
			name: "only two tidb nodes",
			nodes: []*corev1.Node{
				getTestNode(clusterID, "tidb", "us-west-1a"),
				getTestNode(clusterID, "tidb", "us-west-1b"),
				getTestNode("other", "tidb", "us-west-1a"),
				getTestNode("other", "tidb", "us-west-1b"),
				getTestNode("other", "tidb", "us-west-1c"),
			},
			wantExtraTiDBZone: "us-west-1c",
		},
		{
			name: "three tidb nodes",
			nodes: []*corev1.Node{
				getTestNode(clusterID, "tidb", "us-west-1a"),
				getTestNode(clusterID, "tidb", "us-west-1b"),
				getTestNode(clusterID, "tidb", "us-west-1c"),
			},
			wantExtraTiDBZone: "us-west-1a",
		},
		{
			name: "four tidb nodes",
			nodes: []*corev1.Node{
				getTestNode(clusterID, "tidb", "us-west-1a"),
				getTestNode(clusterID, "tidb", "us-west-1b"),
				getTestNode(clusterID, "tidb", "us-west-1c"),
				getTestNode(clusterID, "tidb", "us-west-1a"),
			},
			wantExtraTiDBZone: "us-west-1b",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			var objects []client.Object
			for _, node := range tt.nodes {
				objects = append(objects, node)
			}
			cli := fake.NewClientBuilder().WithScheme(scheme.Scheme).WithObjects(objects...).Build()

			gotExtraTiDBZone, err := GetExtraTiDBZone(ctx, cli, clusterID, zones)
			if (err != nil) != tt.wantErr {
				t.Errorf("GetExtraTiDBZone() error = %v, wantErr %v", err, tt.wantErr)
				return
			}
			if gotExtraTiDBZone != tt.wantExtraTiDBZone {
				t.Errorf("GetExtraTiDBZone() gotExtraTiDBZone = %v, want %v", gotExtraTiDBZone, tt.wantExtraTiDBZone)
			}
		})
	}
}

func getTestTC(slots string, tidbReplicas int32) *opv1alpha1.TidbCluster {
	return &opv1alpha1.TidbCluster{
		ObjectMeta: metav1.ObjectMeta{
			Annotations: map[string]string{
				oplabel.AnnTiDBDeleteSlots: slots,
			},
		},
		Spec: opv1alpha1.TidbClusterSpec{
			TiDB: &opv1alpha1.TiDBSpec{
				Replicas: tidbReplicas,
			},
		},
	}
}

func TestMutateTCToScaleInExtraTiDBPod(t *testing.T) {
	tests := []struct {
		name         string
		tc           *opv1alpha1.TidbCluster
		podName      string
		tidbReplicas int32
		wantErr      bool
		wantSlots    string
	}{
		{
			name:         "scale in extra tidb pod",
			tc:           getTestTC("[1,2,3]", 2),
			podName:      "db-tidb-11",
			tidbReplicas: 1,
			wantErr:      false,
			wantSlots:    "[1,2,3,11]",
		},
		{
			name:         "no delete slots",
			tc:           getTestTC("", 2),
			podName:      "db-tidb-1",
			tidbReplicas: 1,
			wantErr:      false,
			wantSlots:    "[1]",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if err := MutateTCToScaleInExtraTiDBPod(tt.tc, tt.podName, tt.tidbReplicas); (err != nil) != tt.wantErr {
				t.Errorf("MutateTCToScaleInExtraTiDBPod() error = %v, wantErr %v", err, tt.wantErr)
				return
			}

			if tt.tc.Annotations[oplabel.AnnTiDBDeleteSlots] != tt.wantSlots {
				t.Errorf("MutateTCToScaleInExtraTiDBPod() gotSlots = %v, want %v", tt.tc.Annotations[oplabel.AnnTiDBDeleteSlots], tt.wantSlots)
			}

			if tt.tc.Spec.TiDB.Replicas != tt.tidbReplicas {
				t.Errorf("MutateTCToScaleInExtraTiDBPod() gotReplicas = %v, want %v", tt.tc.Spec.TiDB.Replicas, tt.tidbReplicas)
			}
		})
	}
}