package server

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"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func TestEnsureTiProxyGracefulRestartConfigured(t *testing.T) {
	tests := []struct {
		name     string
		provider apiv1.ProviderType
		lbPort   int32
	}{
		{name: "aws", provider: apiv1.ProviderType_AWS},
		{name: "alicloud", provider: apiv1.ProviderType_ALICLOUD},
		{name: "azure uses service port for annotations", provider: apiv1.ProviderType_AZURE, lbPort: 4000},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			ctx := context.Background()
			sharedPoolID := "200"
			dataplaneName := "dp-1"
			tiproxyComponent := "100"
			namespace := entity.GetDataPlaneNamespace(sharedPoolID)
			tiproxyGroupName := tiProxyGroupName(tiproxyComponent)

			scheme := runtime.NewScheme()
			require.NoError(t, v1alpha1.AddToScheme(scheme))
			tiproxyGroup := &v1alpha1.TiProxyGroup{
				ObjectMeta: metav1.ObjectMeta{
					Name:      tiproxyGroupName,
					Namespace: namespace,
				},
			}
			remoteCli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tiproxyGroup).Build()

			sharedPool := &entity.SharedPool{SharedPool: apiv1.SharedPool{
				Id:            sharedPoolID,
				DataplaneName: dataplaneName,
				Provider:      &apiv1.Provider{Type: tt.provider},
				ComponentGroups: []*apiv1.Component{
					{Id: tiproxyComponent, ComponentType: apiv1.ComponentType_TIPROXY},
				},
			}}
			tidbGroup := &entity.TiDBGroup{TiDBGroup: apiv1.TiDBGroup{
				Id:            "tidb-1",
				ClusterId:     tiproxyComponent,
				SharedPoolId:  sharedPoolID,
				DataplaneName: dataplaneName,
				SqlEndpoint: &apiv1.TiDBGroup_SqlEndpoint{
					PrivateLb: &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{Enabled: true},
					PublicLb:  &apiv1.TiDBGroup_SqlEndpoint_LoadBalancer{Enabled: true},
				},
			}}

			tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
			dataplaneClient := dataplanemgrmocks.NewClient(t)
			s := &serverImpl{
				tidbGroupRepo:   tidbGroupRepo,
				dataplaneClient: dataplaneClient,
			}

			tidbGroupRepo.EXPECT().
				List(mock.Anything, []string{sharedPoolID}, mock.Anything, []string{tiproxyComponent}, mock.Anything, false).
				Return([]*entity.TiDBGroup{tidbGroup}, nil).
				Once()

			var expectedAnnotations map[string]string
			if tt.lbPort > 0 {
				expectedAnnotations = entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, tt.lbPort)
			} else {
				expectedAnnotations = entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool)
			}

			for _, lbName := range []string{entity.GetPrivateSvcName(tidbGroup.GetId()), entity.GetPublicSvcName(tidbGroup.GetId())} {
				lbName := lbName
				lbReturn := &dplanemgrv1.LoadBalancer{
					Id:          lbName,
					Annotations: map[string]string{"existing": "kept"},
				}
				if tt.lbPort > 0 {
					lbReturn.Port = tt.lbPort
				}
				dataplaneClient.EXPECT().
					GetLoadBalancer(mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetLoadBalancerRequest) bool {
						return req.GetDataPlaneId() == dataplaneName && req.GetId() == lbName
					})).
					Return(lbReturn, nil).
					Once()
				dataplaneClient.EXPECT().
					UpdateLoadBalancer(mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.UpdateLoadBalancerRequest) bool {
						lb := req.GetLoadbalancer()
						if req.GetDataPlaneId() != dataplaneName || req.GetId() != lbName || !lb.GetGracefulUnhealthyEnabled() {
							return false
						}
						if lb.GetAnnotations()["existing"] != "kept" {
							return false
						}
						for key, value := range expectedAnnotations {
							if lb.GetAnnotations()[key] != value {
								return false
							}
						}
						return true
					})).
					Return(&dplanemgrv1.LoadBalancer{}, nil).
					Once()
			}

			err := s.ensureTiProxyGracefulRestartConfigured(ctx, &tiProxyBackendActionContext{
				sharedPool:         sharedPool,
				remoteCli:          remoteCli,
				tiproxyComponent:   tiproxyComponent,
				tiproxyGroupName:   tiproxyGroupName,
				dataPlaneNamespace: namespace,
			})
			require.NoError(t, err)

			var updatedTiProxyGroup v1alpha1.TiProxyGroup
			require.NoError(t, remoteCli.Get(ctx, types.NamespacedName{Namespace: namespace, Name: tiproxyGroupName}, &updatedTiProxyGroup))
			require.Equal(t, "86400", updatedTiProxyGroup.Spec.Template.Annotations[entity.TiProxyGracefulShutdownDeleteDelaySecondsAnnotation])
		})
	}
}