package gke

import (
	"context"
	"testing"

	opv1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"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"

	. "github.com/tidbcloud/eks-provider/controllers/upgradeoperation/common"
	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	"github.com/tidbcloud/infra-provider/pkg/consts"
	"github.com/tidbcloud/infra-provider/pkg/gkemindomains"
	"github.com/tidbcloud/infra-provider/pkg/namer"
)

func TestResolveMinDomainsUpgradeDecision(t *testing.T) {
	tests := []struct {
		name            string
		unlabeled       bool
		initialDecision gkemindomains.UpgradeDecision
		wantDecision    gkemindomains.UpgradeDecision
		wantPersisted   bool
	}{
		{
			name:          "enable minDomains for a new upgrade",
			wantDecision:  gkemindomains.UpgradeDecisionMinDomains,
			wantPersisted: true,
		},
		{
			name:          "preserve an in-progress unlabeled upgrade",
			unlabeled:     true,
			wantDecision:  gkemindomains.UpgradeDecisionWithoutMinDomains,
			wantPersisted: true,
		},
		{
			name:            "do not recalculate a persisted decision",
			unlabeled:       true,
			initialDecision: gkemindomains.UpgradeDecisionMinDomains,
			wantDecision:    gkemindomains.UpgradeDecisionMinDomains,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			cluster := &cloudv1.Cluster{ObjectMeta: metav1.ObjectMeta{
				Name: "cluster", Namespace: "namespace",
			}}
			uo := &cloudv1.UpgradeOperation{
				ObjectMeta: metav1.ObjectMeta{Name: "upgrade", Namespace: cluster.Namespace},
				Status: cloudv1.UpgradeOperationStatus{CurrentStep: &cloudv1.Step{
					Target: cloudv1.TiDBNodes,
				}},
			}
			if tt.initialDecision != "" {
				uo.Status.CurrentStep.Options = map[string]string{
					gkemindomains.UpgradeDecisionOption: string(tt.initialDecision),
				}
			}

			infraScheme := runtime.NewScheme()
			require.NoError(t, cloudv1.AddToScheme(infraScheme))
			infraCli := fake.NewClientBuilder().
				WithScheme(infraScheme).
				WithStatusSubresource(&cloudv1.UpgradeOperation{}).
				WithObjects(uo).
				Build()

			dataplaneScheme := runtime.NewScheme()
			require.NoError(t, appsv1.AddToScheme(dataplaneScheme))
			objects := []client.Object{}
			if tt.unlabeled {
				objects = append(objects, &appsv1.Deployment{ObjectMeta: metav1.ObjectMeta{
					Name:      namer.BuildOverProvisioningName("pd-1-29"),
					Namespace: cluster.PrefixedId(),
					Labels: map[string]string{
						consts.InstanceLabelKey:  consts.ClusterPrefix,
						consts.ComponentLabelKey: consts.OverProvisioningComponentLabelValue,
					},
				}})
			}
			gkeCli := fake.NewClientBuilder().
				WithScheme(dataplaneScheme).
				WithObjects(objects...).
				Build()

			u := &upgrader{infraCli: infraCli}
			decision, persisted, err := u.resolveMinDomainsUpgradeDecision(ctx, gkeCli, uo, cluster)

			require.NoError(t, err)
			require.Equal(t, tt.wantDecision, decision)
			require.Equal(t, tt.wantPersisted, persisted)
			persistedUO := &cloudv1.UpgradeOperation{}
			require.NoError(t, infraCli.Get(ctx, client.ObjectKeyFromObject(uo), persistedUO))
			require.Equal(t, string(tt.wantDecision),
				persistedUO.GetCurrentStepOption(gkemindomains.UpgradeDecisionOption))
		})
	}
}

func TestMutateTCForGKEUpgradeRestartRecoversMissingMinDomains(t *testing.T) {
	const desiredVersion = "1.30"
	tc := &opv1alpha1.TidbCluster{Spec: opv1alpha1.TidbClusterSpec{
		TiKV: &opv1alpha1.TiKVSpec{
			Replicas: 3,
			ComponentSpec: opv1alpha1.ComponentSpec{Annotations: map[string]string{
				RestartTiDBAnnotation: desiredVersion,
			}},
		},
	}}
	withoutMinDomains := tc.DeepCopy()

	require.True(t, mutateTCForGKEUpgradeRestart(
		tc, desiredVersion, []cloudv1.ComponentType{cloudv1.ComponentTiKV}, true,
	))
	require.Len(t, tc.Spec.TiKV.TopologySpreadConstraints, 1)
	require.Equal(t, int32(gkemindomains.RequiredZoneCount), *tc.Spec.TiKV.TopologySpreadConstraints[0].MinDomains)
	require.Equal(t, desiredVersion, tc.Spec.TiKV.NodeSelector[consts.K8sVersionLabelKey])
	require.False(t, mutateTCForGKEUpgradeRestart(
		tc, desiredVersion, []cloudv1.ComponentType{cloudv1.ComponentTiKV}, true,
	))

	require.False(t, mutateTCForGKEUpgradeRestart(
		withoutMinDomains, desiredVersion, []cloudv1.ComponentType{cloudv1.ComponentTiKV}, false,
	))
	require.Empty(t, withoutMinDomains.Spec.TiKV.TopologySpreadConstraints)
}