//go:build unit || all
// +build unit all

package common

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/eks-provider/controllers/cluster/task"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gkeheterogeneous"
	"github.com/tidbcloud/infra-provider/pkg/namer"
	scheme2 "github.com/tidbcloud/infra-provider/pkg/scheme"
)

func Test_distTaskSupported(t *testing.T) {
	type args struct {
		version string
	}
	tests := []struct {
		name string
		args args
		want bool
	}{
		{name: "v7.5.4", args: args{version: "v7.5.4"}, want: false},
		{name: "v8.1", args: args{version: "v8.1"}, want: true},
		{name: "v8.1.1", args: args{version: "v8.1.1"}, want: true},
		{name: "v8.2", args: args{version: "v8.2"}, want: true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := distTaskSupported(tt.args.version); got != tt.want {
				t.Errorf("distTaskSupported() = %v, want %v", got, tt.want)
			}
		})
	}
}

func TestTiDBGroupEnsureZeroReplicasExtraGroupCleanedDeletesGKEHeterogeneousIntent(t *testing.T) {
	ctx := context.Background()
	cluster, extraGroup := testTiDBGroupCleanupCluster(cloudv1.ProviderTypeGKE, "group1")
	intentKey := cluster.TiDBGKEComputeClassBindingsAnnoKey(extraGroup.Name)
	updateTestGKEComputeClassBinding(
		t,
		cluster,
		intentKey,
		"tidb-group1-priority",
		[]string{"n2-standard-8", "n2-standard-16"},
		testGKEComputeClassNodePoolNames(cluster, "n2-standard-8", "n2-standard-16"),
	)
	primary := testTiDBGroupGatekeeperAssign(t, cluster.Name, extraGroup.Name, "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
	companion := testTiDBGroupGatekeeperAssign(t, cluster.Name, extraGroup.Name, "topology-spread-assign", gkeheterogeneous.GatekeeperAssignPurposeTopologySpread)
	remoteCli := fake.NewClientBuilder().WithScheme(scheme2.Scheme).WithRuntimeObjects(primary, companion).Build()
	tidbGroup := testTiDBGroupCleanupTask(cluster, remoteCli)

	err := tidbGroup.ensureZeroReplicasExtraGroupCleaned(ctx, extraGroup)
	require.NoError(t, err)
	require.NotContains(t, cluster.GetAnnotations(), intentKey)

	gotPrimary := testTiDBGroupGatekeeperAssignLookup()
	err = remoteCli.Get(ctx, client.ObjectKey{Name: primary.GetName()}, gotPrimary)
	require.True(t, apierrors.IsNotFound(err))
	gotCompanion := testTiDBGroupGatekeeperAssignLookup()
	err = remoteCli.Get(ctx, client.ObjectKey{Name: companion.GetName()}, gotCompanion)
	require.True(t, apierrors.IsNotFound(err))
}

func TestTiDBGroupEnsureZeroReplicasExtraGroupCleanedSkipsWithoutGKEHeterogeneousIntent(t *testing.T) {
	tests := []struct {
		name       string
		provider   string
		withIntent bool
	}{
		{
			name:     "GKE without intent",
			provider: cloudv1.ProviderTypeGKE,
		},
		{
			name:       "non GKE with intent",
			provider:   cloudv1.ProviderTypeEKS,
			withIntent: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			cluster, extraGroup := testTiDBGroupCleanupCluster(tt.provider, "group1")
			intentKey := cluster.TiDBGKEComputeClassBindingsAnnoKey(extraGroup.Name)
			if tt.withIntent {
				updateTestGKEComputeClassBinding(
					t,
					cluster,
					intentKey,
					"tidb-group1-priority",
					[]string{"n2-standard-8", "n2-standard-16"},
					testGKEComputeClassNodePoolNames(cluster, "n2-standard-8", "n2-standard-16"),
				)
			}
			primary := testTiDBGroupGatekeeperAssign(t, cluster.Name, extraGroup.Name, "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
			remoteCli := fake.NewClientBuilder().WithScheme(scheme2.Scheme).WithRuntimeObjects(primary).Build()
			tidbGroup := testTiDBGroupCleanupTask(cluster, remoteCli)

			err := tidbGroup.ensureZeroReplicasExtraGroupCleaned(ctx, extraGroup)
			require.NoError(t, err)

			gotPrimary := testTiDBGroupGatekeeperAssignLookup()
			require.NoError(t, remoteCli.Get(ctx, client.ObjectKey{Name: primary.GetName()}, gotPrimary))
			if tt.withIntent {
				require.Contains(t, cluster.GetAnnotations(), intentKey)
			} else {
				require.NotContains(t, cluster.GetAnnotations(), intentKey)
			}
		})
	}
}

func TestTiDBGroupEnsureZeroReplicasExtraGroupCleanedStopsWhenGatekeeperDeleteFails(t *testing.T) {
	ctx := context.Background()
	cluster, extraGroup := testTiDBGroupCleanupCluster(cloudv1.ProviderTypeGKE, "group1")
	extraGroup.Status.TiDB = &cloudv1.TiDBStatus{Status: cloudv1.ComponentNormal}
	intentKey := cluster.TiDBGKEComputeClassBindingsAnnoKey(extraGroup.Name)
	updateTestGKEComputeClassBinding(
		t,
		cluster,
		intentKey,
		"tidb-group1-priority",
		[]string{"n2-standard-8", "n2-standard-16"},
		testGKEComputeClassNodePoolNames(cluster, "n2-standard-8", "n2-standard-16"),
	)
	primary := testTiDBGroupGatekeeperAssign(t, cluster.Name, extraGroup.Name, "compute-class-assign", gkeheterogeneous.GatekeeperAssignPurposeNodeSelector)
	labels := primary.GetLabels()
	labels[consts.ClusterLabelKey] = "other-cluster"
	primary.SetLabels(labels)
	remoteCli := fake.NewClientBuilder().WithScheme(scheme2.Scheme).WithRuntimeObjects(primary).Build()
	tidbGroup := testTiDBGroupCleanupTask(cluster, remoteCli)

	err := tidbGroup.ensureZeroReplicasExtraGroupCleaned(ctx, extraGroup)
	require.ErrorContains(t, err, "delete GKE heterogeneous gatekeeper assign")
	require.Contains(t, cluster.GetAnnotations(), intentKey)
	require.NotNil(t, extraGroup.Status.TiDB)
}

func testTiDBGroupCleanupCluster(providerType, groupID string) (*cloudv1.Cluster, *cloudv1.ExtraGroupCluster) {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Name:        "123",
			Namespace:   "default",
			Labels:      map[string]string{cloudv1.ProviderTypeLabel: providerType},
			Annotations: map[string]string{},
		},
		Spec: cloudv1.ClusterSpec{
			ExtraGroups: []cloudv1.ExtraGroupSpec{
				{
					Name: groupID,
					TiDB: cloudv1.TiDB{
						Replicas:     0,
						InstanceType: "n2-standard-8",
					},
				},
			},
		},
		Status: cloudv1.ClusterStatus{
			ExtraEndpoints: []cloudv1.ExtraGroupStatus{
				{
					Name: groupID,
				},
			},
		},
	}
	extraGroup, _ := cluster.GetExtraGroup(groupID)
	return cluster, extraGroup
}

func testTiDBGroupCleanupTask(cluster *cloudv1.Cluster, remoteCli client.Client) *TiDBGroup {
	return &TiDBGroup{
		BaseTask: &task.BaseTask{
			Context: &task.Context{
				Cluster: cluster,
				Namer: namer.NewNamer().
					WithClusterName(cluster.Name).
					WithTiDBNamespace(cluster.PrefixedId()),
				EnvProvider: &testEnvProvider{client: remoteCli},
			},
		},
	}
}

func testTiDBGroupGatekeeperAssign(t *testing.T, clusterID, tidbGroupID, suffix, purpose string) *unstructured.Unstructured {
	t.Helper()
	obj := &unstructured.Unstructured{
		Object: map[string]interface{}{
			"apiVersion": "mutations.gatekeeper.sh/v1",
			"kind":       "Assign",
			"metadata": map[string]interface{}{
				"name": gkeheterogeneous.ResourceName(clusterID, cloudv1.ComponentTiDB, tidbGroupID, suffix),
			},
			"spec": map[string]interface{}{
				"match": map[string]interface{}{
					"scope":      "Namespaced",
					"namespaces": []interface{}{"tidb" + clusterID},
				},
			},
		},
	}
	labels := map[string]string{
		consts.ManagedByLabelKey:                         gkeheterogeneous.ManagedByClusterAgentLabelValue,
		consts.ClusterLabelKey:                           clusterID,
		consts.LabelKeyTiDBGroupId:                       tidbGroupID,
		gkeheterogeneous.HeterogeneousDeployLabelKey:     gkeheterogeneous.GatekeeperAssignDeployLabelValue,
		gkeheterogeneous.HeterogeneousComponentLabelKey:  string(cloudv1.ComponentTiDB),
		gkeheterogeneous.GatekeeperAssignPurposeLabelKey: purpose,
	}
	obj.SetLabels(labels)
	if purpose == gkeheterogeneous.GatekeeperAssignPurposeTopologySpread {
		require.NoError(t, unstructured.SetNestedField(obj.Object, "spec.topologySpreadConstraints[topologyKey:'topology.kubernetes.io/zone'].nodeAffinityPolicy", "spec", "location"))
		require.NoError(t, unstructured.SetNestedField(obj.Object, "Ignore", "spec", "parameters", "assign", "value"))
	}
	return obj
}

func testTiDBGroupGatekeeperAssignLookup() *unstructured.Unstructured {
	obj := &unstructured.Unstructured{}
	obj.SetAPIVersion("mutations.gatekeeper.sh/v1")
	obj.SetKind("Assign")
	return obj
}