package server

import (
	"context"
	"fmt"

	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/pkg/k8s"
	"github.com/tidbcloud/tiinf/monitor/tilog"
	"go.uber.org/zap"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func (s *serverImpl) UpdateSharedPoolVcpusLimit(ctx context.Context, sharedPoolId string, vcpusLimit uint64) (*entity.SharedPool, error) {
	if sharedPoolId == "" {
		return nil, status.Errorf(codes.InvalidArgument, "sharedPool id is required")
	}

	var updatedSharedPool *entity.SharedPool
	if err := s.transactioner.Transaction(ctx, func(ctx context.Context) error {
		unlockFunc, err := s.lock(ctx, sharedPoolId)
		if err != nil {
			tilog.Error(ctx, "failed to lock sharedPool", zap.String("sharedPoolId", sharedPoolId), zap.Error(err))
			return err
		}
		defer func() {
			if unlockErr := unlockFunc(); unlockErr != nil {
				tilog.Error(ctx, "failed to unlock sharedPool", zap.String("sharedPoolId", sharedPoolId), zap.Error(unlockErr))
			}
		}()

		sharedPool, err := s.sharedPoolRepo.Load(ctx, sharedPoolId, false)
		if err != nil {
			if status.Code(err) != codes.NotFound {
				return status.Errorf(codes.Internal, "failed to load sharedPool: %s", err)
			}
			return status.Errorf(codes.NotFound, "sharedPool %s not found", sharedPoolId)
		}
		if err := guardSharedPoolMutation(sharedPool, sharedPoolMutationUpdateVCPUsLimit); err != nil {
			return err
		}

		remoteCli, err := s.dataplaneManager.GetClient(sharedPool.GetDataplaneName())
		if err != nil {
			return status.Errorf(codes.Internal, "failed to get dataplane client: %s", err)
		}

		oldVcpusLimit := sharedPool.GetVcpusLimit()

		revertPolicies := func(revertVcpusLimit uint64) {
			if revertErr := k8s.EnsureSharedPoolLimitRange(ctx, remoteCli, sharedPool.GetDataPlaneNamespace(), revertVcpusLimit); revertErr != nil {
				tilog.Error(ctx, "failed to revert limit range after sharedPool policy update failure",
					zap.String("sharedPoolId", sharedPoolId),
					zap.Uint64("revertVcpusLimit", revertVcpusLimit),
					zap.Error(revertErr))
			}
			if revertErr := k8s.EnsureResourceQuota(ctx, remoteCli, sharedPool.GetDataPlaneNamespace(), revertVcpusLimit); revertErr != nil {
				tilog.Error(ctx, "failed to revert resource quota after sharedPool policy update failure",
					zap.String("sharedPoolId", sharedPoolId),
					zap.Uint64("revertVcpusLimit", revertVcpusLimit),
					zap.Error(revertErr))
			}
		}

		if err := k8s.EnsureSharedPoolLimitRange(ctx, remoteCli, sharedPool.GetDataPlaneNamespace(), vcpusLimit); err != nil {
			return status.Errorf(codes.Internal, "failed to ensure limit range: %s", err)
		}

		if err := k8s.EnsureResourceQuota(ctx, remoteCli, sharedPool.GetDataPlaneNamespace(), vcpusLimit); err != nil {
			revertPolicies(oldVcpusLimit)
			return status.Errorf(codes.Internal, "failed to ensure resource quota: %s", err)
		}

		sharedPool.VcpusLimit = vcpusLimit
		if err := s.sharedPoolRepo.Save(ctx, sharedPool); err != nil {
			revertPolicies(oldVcpusLimit)
			tilog.Error(ctx, "failed to save sharedPool meta", zap.Error(err))
			return fmt.Errorf("failed to save sharedPool meta: %w", err)
		}

		updatedSharedPool = sharedPool
		return nil
	}); err != nil {
		return nil, err
	}

	return updatedSharedPool, nil
}