package server

import (
	"context"
	"testing"

	"github.com/pingcap/tidb-operator/api/v2/core/v1alpha1"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	clusteragentv1 "github.com/tidbcloud/cluster-service-ng/api/tidbcloud/cluster-agent/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	dataplanemocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	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 TestPrepareTiProxyGracefulRestartForTargetsBlocksPausedPoolBeforeKubernetesMutation(t *testing.T) {
	ctx := context.Background()
	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "sp-1",
			DataplaneName: "dp-1",
			ComponentGroups: []*entity.Component{
				{Id: "tiproxy-1", ComponentType: apiv1.ComponentType_TIPROXY},
			},
		},
		PauseResumeState:      entity.PauseResumeStatePaused,
		SuspendComputeDesired: true,
	}

	scheme := runtime.NewScheme()
	require.NoError(t, v1alpha1.AddToScheme(scheme))
	tiproxyGroup := &v1alpha1.TiProxyGroup{ObjectMeta: metav1.ObjectMeta{
		Name:      tiProxyGroupName("tiproxy-1"),
		Namespace: sharedPool.GetDataPlaneNamespace(),
	}}
	remoteCli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tiproxyGroup).Build()
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	dataplaneManager := dataplanemocks.NewManager(t)
	sharedPoolRepo.EXPECT().Load(mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil).Once()
	dataplaneManager.EXPECT().GetClient(sharedPool.GetDataplaneName()).Return(remoteCli, nil).Once()

	s := &serverImpl{sharedPoolRepo: sharedPoolRepo, dataplaneManager: dataplaneManager}
	_, err := s.prepareTiProxyGracefulRestartForTargets(ctx, sharedPool.GetId(), []*entity.TargetComponent{{
		Id:            "tiproxy-1",
		ComponentType: apiv1.ComponentType_TIPROXY,
	}}, apiv1.SetSharedPoolComponentConfigVersionRequest_ROLLING_UPDATE)

	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	var updated v1alpha1.TiProxyGroup
	require.NoError(t, remoteCli.Get(ctx, types.NamespacedName{Name: tiproxyGroup.Name, Namespace: tiproxyGroup.Namespace}, &updated))
	assert.Empty(t, updated.Spec.Template.Annotations)
}

func TestTiProxyPodsNeedingLegacyNoNewConnections(t *testing.T) {
	resp := &clusteragentv1.ProbeTiProxyHealthAPIResp{Results: []*clusteragentv1.TiProxyHealthAPIProbeResult{
		{PodName: "tiproxy-a", Status: clusteragentv1.TiProxyHealthAPIProbeStatus_SUPPORTED},
		{PodName: "tiproxy-b", Status: clusteragentv1.TiProxyHealthAPIProbeStatus_UNSUPPORTED},
		{PodName: "tiproxy-c", Status: clusteragentv1.TiProxyHealthAPIProbeStatus_UNSUPPORTED},
	}}

	pods, err := tiProxyPodsNeedingLegacyNoNewConnections(resp, []string{"tiproxy-a", "tiproxy-b", "tiproxy-c"})

	require.NoError(t, err)
	assert.Equal(t, []string{"tiproxy-b", "tiproxy-c"}, pods)
}

func TestTiProxyPodsNeedingLegacyNoNewConnectionsFailsOnUnknown(t *testing.T) {
	resp := &clusteragentv1.ProbeTiProxyHealthAPIResp{Results: []*clusteragentv1.TiProxyHealthAPIProbeResult{
		{PodName: "tiproxy-a", Status: clusteragentv1.TiProxyHealthAPIProbeStatus_UNKNOWN, Message: "timeout"},
	}}

	_, err := tiProxyPodsNeedingLegacyNoNewConnections(resp, []string{"tiproxy-a"})

	require.Error(t, err)
	assert.Equal(t, codes.FailedPrecondition, status.Code(err))
	assert.Contains(t, err.Error(), "timeout")
}