package server

import (
	"context"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/pkg/dataplane/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/k8s"
	repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	corev1 "k8s.io/api/core/v1"
	apierrors "k8s.io/apimachinery/pkg/api/errors"
	"k8s.io/apimachinery/pkg/api/resource"
	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 TestUpdateSharedPoolVcpusLimit(t *testing.T) {
	tests := []struct {
		name        string
		vcpusLimit  uint64
		setupClient func(*entity.SharedPool) (*runtime.Scheme, types.NamespacedName, types.NamespacedName, *corev1.ResourceQuota, *corev1.LimitRange)
	}{
		{
			name:       "set resource quota",
			vcpusLimit: 106,
			setupClient: func(sharedPool *entity.SharedPool) (*runtime.Scheme, types.NamespacedName, types.NamespacedName, *corev1.ResourceQuota, *corev1.LimitRange) {
				scheme := runtime.NewScheme()
				require.NoError(t, corev1.AddToScheme(scheme))
				return scheme,
					types.NamespacedName{Name: k8s.SharedPoolVCPUCeilingResourceQuotaName, Namespace: sharedPool.GetDataPlaneNamespace()},
					types.NamespacedName{Name: k8s.SharedPoolDefaultCPURequestLimitRangeName, Namespace: sharedPool.GetDataPlaneNamespace()},
					nil,
					nil
			},
		},
		{
			name:       "remove resource quota when limit is zero",
			vcpusLimit: 0,
			setupClient: func(sharedPool *entity.SharedPool) (*runtime.Scheme, types.NamespacedName, types.NamespacedName, *corev1.ResourceQuota, *corev1.LimitRange) {
				scheme := runtime.NewScheme()
				require.NoError(t, corev1.AddToScheme(scheme))
				rq := &corev1.ResourceQuota{
					ObjectMeta: metav1.ObjectMeta{
						Name:      k8s.SharedPoolVCPUCeilingResourceQuotaName,
						Namespace: sharedPool.GetDataPlaneNamespace(),
					},
				}
				lr := &corev1.LimitRange{
					ObjectMeta: metav1.ObjectMeta{
						Name:      k8s.SharedPoolDefaultCPURequestLimitRangeName,
						Namespace: sharedPool.GetDataPlaneNamespace(),
					},
					Spec: corev1.LimitRangeSpec{
						Limits: []corev1.LimitRangeItem{{
							Type: corev1.LimitTypeContainer,
							DefaultRequest: corev1.ResourceList{
								corev1.ResourceCPU: resource.MustParse("20m"),
							},
						}},
					},
				}
				return scheme,
					types.NamespacedName{Name: rq.Name, Namespace: rq.Namespace},
					types.NamespacedName{Name: lr.Name, Namespace: lr.Namespace},
					rq,
					lr
			},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			lockRepo := repomnocks.NewLockRepo(t)
			dataplaneManager := mocks.NewManager(t)

			sharedPool := createTestSharedPool()
			scheme, quotaKey, limitRangeKey, existingQuota, existingLimitRange := tt.setupClient(sharedPool)
			builder := fake.NewClientBuilder().WithScheme(scheme)
			if existingQuota != nil {
				builder = builder.WithObjects(existingQuota)
			}
			if existingLimitRange != nil {
				builder = builder.WithObjects(existingLimitRange)
			}
			k8sClient := builder.Build()

			sharedPoolRepo.On("Load", mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil)
			sharedPoolRepo.On("Save", mock.Anything, mock.MatchedBy(func(sp *entity.SharedPool) bool {
				return sp.GetId() == sharedPool.GetId() && sp.GetVcpusLimit() == tt.vcpusLimit
			})).Return(nil)
			lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
				return lock.Key == sharedPool.GetId()
			})).Return(func() error { return nil }, nil)
			dataplaneManager.On("GetClient", sharedPool.GetDataplaneName()).Return(k8sClient, nil)

			server := &serverImpl{
				transactioner:    &mockTransactioner{},
				sharedPoolRepo:   sharedPoolRepo,
				lockRepo:         lockRepo,
				dataplaneManager: dataplaneManager,
			}

			result, err := server.UpdateSharedPoolVcpusLimit(context.Background(), sharedPool.GetId(), tt.vcpusLimit)
			require.NoError(t, err)
			require.NotNil(t, result)
			assert.Equal(t, tt.vcpusLimit, result.GetVcpusLimit())

			storedQuota := &corev1.ResourceQuota{}
			err = k8sClient.Get(context.Background(), quotaKey, storedQuota)
			storedLimitRange := &corev1.LimitRange{}
			limitRangeErr := k8sClient.Get(context.Background(), limitRangeKey, storedLimitRange)
			if tt.vcpusLimit == 0 {
				require.Error(t, err)
				assert.True(t, apierrors.IsNotFound(err))
				require.Error(t, limitRangeErr)
				assert.True(t, apierrors.IsNotFound(limitRangeErr))
				return
			}

			require.NoError(t, err)
			cpuRequest := storedQuota.Spec.Hard[corev1.ResourceRequestsCPU]
			assert.EqualValues(t, 106, cpuRequest.Value())
			require.NoError(t, limitRangeErr)
			require.Len(t, storedLimitRange.Spec.Limits, 1)
			assert.Equal(t, "20m", storedLimitRange.Spec.Limits[0].DefaultRequest.Cpu().String())
		})
	}
}

func TestUpdateSharedPoolVcpusLimitRequiresSharedPoolID(t *testing.T) {
	server := &serverImpl{}
	_, err := server.UpdateSharedPoolVcpusLimit(context.Background(), "", 1)
	require.Error(t, err)
	assert.Equal(t, codes.InvalidArgument, status.Code(err))
}

func TestUpdateSharedPoolVcpusLimitRevertsPoliciesWhenSaveFails(t *testing.T) {
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	lockRepo := repomnocks.NewLockRepo(t)
	dataplaneManager := mocks.NewManager(t)

	sharedPool := createTestSharedPool()
	sharedPool.VcpusLimit = 80

	scheme := runtime.NewScheme()
	require.NoError(t, corev1.AddToScheme(scheme))
	existingQuota := &corev1.ResourceQuota{
		ObjectMeta: metav1.ObjectMeta{
			Name:      k8s.SharedPoolVCPUCeilingResourceQuotaName,
			Namespace: sharedPool.GetDataPlaneNamespace(),
		},
		Spec: corev1.ResourceQuotaSpec{
			Hard: corev1.ResourceList{
				corev1.ResourceRequestsCPU: resource.MustParse("80"),
			},
		},
	}
	existingLimitRange := &corev1.LimitRange{
		ObjectMeta: metav1.ObjectMeta{
			Name:      k8s.SharedPoolDefaultCPURequestLimitRangeName,
			Namespace: sharedPool.GetDataPlaneNamespace(),
		},
		Spec: corev1.LimitRangeSpec{
			Limits: []corev1.LimitRangeItem{{
				Type: corev1.LimitTypeContainer,
				DefaultRequest: corev1.ResourceList{
					corev1.ResourceCPU: resource.MustParse("20m"),
				},
			}},
		},
	}
	k8sClient := fake.NewClientBuilder().WithScheme(scheme).WithObjects(existingQuota, existingLimitRange).Build()

	sharedPoolRepo.On("Load", mock.Anything, sharedPool.GetId(), false).Return(sharedPool, nil)
	sharedPoolRepo.On("Save", mock.Anything, mock.Anything).Return(assert.AnError)
	lockRepo.On("Lock", mock.Anything, mock.MatchedBy(func(lock *entity.Lock) bool {
		return lock.Key == sharedPool.GetId()
	})).Return(func() error { return nil }, nil)
	dataplaneManager.On("GetClient", sharedPool.GetDataplaneName()).Return(k8sClient, nil)

	server := &serverImpl{
		transactioner:    &mockTransactioner{},
		sharedPoolRepo:   sharedPoolRepo,
		lockRepo:         lockRepo,
		dataplaneManager: dataplaneManager,
	}

	_, err := server.UpdateSharedPoolVcpusLimit(context.Background(), sharedPool.GetId(), 106)
	require.Error(t, err)

	storedQuota := &corev1.ResourceQuota{}
	err = k8sClient.Get(context.Background(), types.NamespacedName{
		Name:      k8s.SharedPoolVCPUCeilingResourceQuotaName,
		Namespace: sharedPool.GetDataPlaneNamespace(),
	}, storedQuota)
	require.NoError(t, err)
	cpuRequest := storedQuota.Spec.Hard[corev1.ResourceRequestsCPU]
	assert.EqualValues(t, 80, cpuRequest.Value())

	storedLimitRange := &corev1.LimitRange{}
	err = k8sClient.Get(context.Background(), types.NamespacedName{
		Name:      k8s.SharedPoolDefaultCPURequestLimitRangeName,
		Namespace: sharedPool.GetDataPlaneNamespace(),
	}, storedLimitRange)
	require.NoError(t, err)
	require.Len(t, storedLimitRange.Spec.Limits, 1)
	assert.Equal(t, "20m", storedLimitRange.Spec.Limits[0].DefaultRequest.Cpu().String())
}