package nodepool

import (
	"context"
	"testing"

	"github.com/stretchr/testify/require"
	appsv1 "k8s.io/api/apps/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"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

func TestCreateOverProvisioningPodsUsesGKEV1Layout(t *testing.T) {
	tests := []struct {
		name          string
		useMinDomains bool
	}{
		{name: "minDomains upgrade", useMinDomains: true},
		{name: "legacy upgrade"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			cluster := &cloudv1.Cluster{
				ObjectMeta: metav1.ObjectMeta{
					Labels:      map[string]string{cloudv1.ProviderTypeLabel: cloudv1.ProviderTypeGKE},
					Annotations: map[string]string{},
				},
				Spec: cloudv1.ClusterSpec{Cluster: cloudv1.ClusterTopology{
					TiKV: cloudv1.TiKV{Replicas: 3, InstanceType: "n2-standard-8"},
				}},
			}
			cluster.EnableFeature(consts.ClusterFeatureGCP3LayerNodePool)

			scheme := runtime.NewScheme()
			require.NoError(t, appsv1.AddToScheme(scheme))
			cli := fake.NewClientBuilder().WithScheme(scheme).Build()
			err := CreateOverProvisioningPods(
				context.Background(), cli, cluster, "tc", "ns", "1.30", nil, tt.useMinDomains,
			)
			require.NoError(t, err)

			deployments := &appsv1.DeploymentList{}
			require.NoError(t, cli.List(context.Background(), deployments, client.InNamespace("ns")))
			require.Len(t, deployments.Items, 1)
			constraint := deployments.Items[0].Spec.Template.Spec.TopologySpreadConstraints[0]
			if tt.useMinDomains {
				require.NotNil(t, constraint.MinDomains)
				require.Equal(t, int32(3), *constraint.MinDomains)
			} else {
				require.Nil(t, constraint.MinDomains)
			}
		})
	}
}

func TestCreateOverProvisioningPodsPreservesExistingDeployment(t *testing.T) {
	cluster := &cloudv1.Cluster{
		ObjectMeta: metav1.ObjectMeta{
			Labels:      map[string]string{cloudv1.ProviderTypeLabel: cloudv1.ProviderTypeGKE},
			Annotations: map[string]string{},
		},
		Spec: cloudv1.ClusterSpec{Cluster: cloudv1.ClusterTopology{
			TiKV: cloudv1.TiKV{Replicas: 3, InstanceType: "n2-standard-8"},
		}},
	}
	cluster.EnableFeature(consts.ClusterFeatureGCP3LayerNodePool)

	scheme := runtime.NewScheme()
	require.NoError(t, appsv1.AddToScheme(scheme))
	cli := fake.NewClientBuilder().WithScheme(scheme).Build()
	ctx := context.Background()

	err := CreateOverProvisioningPods(ctx, cli, cluster, "tc", "ns", "1.30", nil, false)
	require.NoError(t, err)

	deployments := &appsv1.DeploymentList{}
	require.NoError(t, cli.List(ctx, deployments, client.InNamespace("ns")))
	require.Len(t, deployments.Items, 1)
	legacy := &deployments.Items[0]
	err = CreateOverProvisioningPods(ctx, cli, cluster, "tc", "ns", "1.30", nil, true)
	require.NoError(t, err)
	current := &appsv1.Deployment{}
	require.NoError(t, cli.Get(ctx, client.ObjectKeyFromObject(legacy), current))
	require.Nil(t, current.Spec.Template.Spec.TopologySpreadConstraints[0].MinDomains)

	exists, err := HasUnlabeledUpgradeOverProvisioning(ctx, cli, "tc", "ns")
	require.NoError(t, err)
	require.True(t, exists)

	otherExists, err := HasUnlabeledUpgradeOverProvisioning(ctx, cli, "other", "ns")
	require.NoError(t, err)
	require.False(t, otherExists)

	v2 := &appsv1.Deployment{ObjectMeta: metav1.ObjectMeta{
		Name:      namer.BuildOverProvisioningName("tidb-1.31"),
		Namespace: "ns",
		Labels: map[string]string{
			consts.InstanceLabelKey:                  "v2-tc",
			consts.ComponentLabelKey:                 consts.OverProvisioningComponentLabelValue,
			consts.OverProvisioningUsedForLabelKey:   consts.OverProvisioningUsedForTidbcluster,
			consts.OverProvisioningComponentLabelKey: string(cloudv1.ComponentTiDB),
		},
	}}
	require.NoError(t, cli.Create(ctx, v2))
	v2Exists, err := HasUnlabeledUpgradeOverProvisioning(ctx, cli, "v2-tc", "ns")
	require.NoError(t, err)
	require.False(t, v2Exists)
}