package deploy

import (
	"context"
	"testing"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/apis/meta/v1/unstructured"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestBackfillTiProxyGracefulRestartAnnotationsInNamespace(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))

	ns1 := entity.GetDataPlaneNamespace("sp-1")
	ns2 := entity.GetDataPlaneNamespace("sp-2")
	existingAnnotation := "existing"
	migratedGroup := &v1alpha1.TiProxyGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db-tiproxy-1",
			Namespace: ns1,
		},
		Spec: v1alpha1.TiProxyGroupSpec{
			Template: v1alpha1.TiProxyTemplate{
				ObjectMeta: v1alpha1.ObjectMeta{
					Annotations: map[string]string{existingAnnotation: "true"},
				},
			},
		},
	}
	alreadyMigratedGroup := &v1alpha1.TiProxyGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db-tiproxy-2",
			Namespace: ns1,
		},
		Spec: v1alpha1.TiProxyGroupSpec{
			Template: v1alpha1.TiProxyTemplate{
				ObjectMeta: v1alpha1.ObjectMeta{
					Annotations: map[string]string{
						entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation: "86400",
					},
				},
			},
		},
	}
	otherNamespaceGroup := &v1alpha1.TiProxyGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db-tiproxy-other",
			Namespace: ns2,
		},
	}
	remoteCli := fake.NewClientBuilder().
		WithScheme(scheme).
		WithObjects(migratedGroup, alreadyMigratedGroup, otherNamespaceGroup).
		Build()

	updated, err := backfillTiProxyGracefulRestartAnnotationsInNamespace(ctx, remoteCli, ns1)
	require.NoError(t, err)
	require.Equal(t, 1, updated)

	var got v1alpha1.TiProxyGroup
	require.NoError(t, remoteCli.Get(ctx, types.NamespacedName{Namespace: ns1, Name: migratedGroup.Name}, &got))
	require.Equal(t, "true", got.Spec.Template.Annotations[existingAnnotation])
	require.Equal(t, "86400", got.Spec.Template.Annotations[entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation])

	require.NoError(t, remoteCli.Get(ctx, types.NamespacedName{Namespace: ns2, Name: otherNamespaceGroup.Name}, &got))
	require.NotContains(t, got.Spec.Template.Annotations, entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation)

	updated, err = backfillTiProxyGracefulRestartAnnotationsInNamespace(ctx, remoteCli, ns1)
	require.NoError(t, err)
	require.Equal(t, 0, updated)
}

func TestBackfillTiProxyGracefulRestartAnnotationsForSharedPools(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))

	sharedPools := []*entity.SharedPool{
		{SharedPool: apiv1.SharedPool{Id: "sp-1", DataplaneName: "dp-1"}},
		{SharedPool: apiv1.SharedPool{Id: "sp-2", DataplaneName: "dp-2"}},
	}
	clients := map[string]client.Client{}
	for _, sharedPool := range sharedPools {
		clients[sharedPool.GetDataplaneName()] = fake.NewClientBuilder().
			WithScheme(scheme).
			WithObjects(newTiProxyGroupCRD("v2.1.0-beta.5"), &v1alpha1.TiProxyGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "db-" + sharedPool.GetId(),
					Namespace: sharedPool.GetDataPlaneNamespace(),
				},
			}).
			Build()
	}

	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	dataplaneManager := dataplanemocks.NewManager(t)
	sharedPoolRepo.EXPECT().
		List(mock.Anything, &repo.SharedPoolFilter{Deleted: false}).
		Return(sharedPools, nil).
		Once()
	for _, sharedPool := range sharedPools {
		sharedPool := sharedPool
		dataplaneManager.EXPECT().
			GetClient(sharedPool.GetDataplaneName()).
			Return(clients[sharedPool.GetDataplaneName()], nil).
			Once()
	}

	m := &migrator{sharedPoolRepo: sharedPoolRepo, dataplaneManager: dataplaneManager}
	require.NoError(t, backfillTiProxyGracefulRestartAnnotationsForSharedPools(ctx, m))

	for _, sharedPool := range sharedPools {
		var got v1alpha1.TiProxyGroup
		require.NoError(t, clients[sharedPool.GetDataplaneName()].Get(ctx, types.NamespacedName{
			Namespace: sharedPool.GetDataPlaneNamespace(),
			Name:      "db-" + sharedPool.GetId(),
		}, &got))
		require.Equal(t, "86400", got.Spec.Template.Annotations[entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation])
	}
}

func TestBackfillTiProxyGracefulRestartAnnotationsForSharedPoolsSkipsUnsafeOperatorVersions(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))

	sharedPools := []*entity.SharedPool{
		{SharedPool: apiv1.SharedPool{Id: "safe", DataplaneName: "dp-safe"}},
		{SharedPool: apiv1.SharedPool{Id: "old", DataplaneName: "dp-old"}},
		{SharedPool: apiv1.SharedPool{Id: "missing-version", DataplaneName: "dp-missing-version"}},
	}
	clients := map[string]client.Client{
		"dp-safe": fake.NewClientBuilder().
			WithScheme(scheme).
			WithObjects(
				newTiProxyGroupCRD("v2.1.0-beta.5"),
				newTiProxyGroupForMigration(sharedPools[0]),
			).
			Build(),
		"dp-old": fake.NewClientBuilder().
			WithScheme(scheme).
			WithObjects(
				newTiProxyGroupCRD("v2.1.0-beta.4"),
				newTiProxyGroupForMigration(sharedPools[1]),
			).
			Build(),
		"dp-missing-version": fake.NewClientBuilder().
			WithScheme(scheme).
			WithObjects(
				newTiProxyGroupCRD(""),
				newTiProxyGroupForMigration(sharedPools[2]),
			).
			Build(),
	}

	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	dataplaneManager := dataplanemocks.NewManager(t)
	sharedPoolRepo.EXPECT().
		List(mock.Anything, &repo.SharedPoolFilter{Deleted: false}).
		Return(sharedPools, nil).
		Once()
	for _, sharedPool := range sharedPools {
		sharedPool := sharedPool
		dataplaneManager.EXPECT().
			GetClient(sharedPool.GetDataplaneName()).
			Return(clients[sharedPool.GetDataplaneName()], nil).
			Once()
	}

	m := &migrator{sharedPoolRepo: sharedPoolRepo, dataplaneManager: dataplaneManager}
	require.NoError(t, backfillTiProxyGracefulRestartAnnotationsForSharedPools(ctx, m))

	var got v1alpha1.TiProxyGroup
	require.NoError(t, clients["dp-safe"].Get(ctx, types.NamespacedName{
		Namespace: sharedPools[0].GetDataPlaneNamespace(),
		Name:      "db-" + sharedPools[0].GetId(),
	}, &got))
	require.Equal(t, "86400", got.Spec.Template.Annotations[entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation])

	require.NoError(t, clients["dp-old"].Get(ctx, types.NamespacedName{
		Namespace: sharedPools[1].GetDataPlaneNamespace(),
		Name:      "db-" + sharedPools[1].GetId(),
	}, &got))
	require.NotContains(t, got.Spec.Template.Annotations, entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation)

	require.NoError(t, clients["dp-missing-version"].Get(ctx, types.NamespacedName{
		Namespace: sharedPools[2].GetDataPlaneNamespace(),
		Name:      "db-" + sharedPools[2].GetId(),
	}, &got))
	require.NotContains(t, got.Spec.Template.Annotations, entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation)
}

func TestTiDBOperatorCanPatchTiProxyGracefulRestartAnnotations(t *testing.T) {
	ctx := context.Background()
	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))

	cases := []struct {
		name    string
		version string
		safe    bool
		wantErr bool
	}{
		{name: "before safe beta", version: "v2.1.0-beta.4", safe: false},
		{name: "safe beta", version: "v2.1.0-beta.5", safe: true},
		{name: "safe rc", version: "v2.1.0-rc.0", safe: true},
		{name: "safe final", version: "v2.1.0", safe: true},
		{name: "safe next minor alpha", version: "v2.2.0-alpha.0", safe: true},
		{name: "missing version", wantErr: true},
		{name: "invalid version", version: "invalid", wantErr: true},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			remoteCli := fake.NewClientBuilder().
				WithScheme(scheme).
				WithObjects(newTiProxyGroupCRD(tc.version)).
				Build()

			safe, _, err := tidbOperatorCanPatchTiProxyGracefulRestartAnnotations(ctx, remoteCli)
			if tc.wantErr {
				require.Error(t, err)
				return
			}
			require.NoError(t, err)
			require.Equal(t, tc.safe, safe)
		})
	}
}

func newTiProxyGroupCRD(version string) *unstructured.Unstructured {
	crd := &unstructured.Unstructured{}
	crd.SetGroupVersionKind(tiProxyGroupCRDGKV)
	crd.SetName(tiProxyGroupCRDName)
	if version != "" {
		crd.SetAnnotations(map[string]string{tidbOperatorVersionAnnotation: version})
	}
	return crd
}

func newTiProxyGroupForMigration(sharedPool *entity.SharedPool) *v1alpha1.TiProxyGroup {
	return &v1alpha1.TiProxyGroup{
		ObjectMeta: metav1.ObjectMeta{
			Name:      "db-" + sharedPool.GetId(),
			Namespace: sharedPool.GetDataPlaneNamespace(),
		},
	}
}