package qos

import (
	"context"

	"git.frostfs.info/TrueCloudLab/frostfs-qos/limiting"
	"git.frostfs.info/TrueCloudLab/frostfs-qos/tagging"
	apistatus "git.frostfs.info/TrueCloudLab/frostfs-sdk-go/client/status"
	"google.golang.org/grpc"
)

func NewMaxActiveRPCLimiterUnaryServerInterceptor(lr *limiting.Limiter) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
		if tag, ok := tagging.IOTagFromContext(ctx); ok && tag == IOTagCritical.String() {
			return handler(ctx, req)
		}

		release, ok := lr.TryAcquire(info.FullMethod)
		if !ok {
			return nil, new(apistatus.ResourceExhausted)
		}
		defer release()

		return handler(ctx, req)
	}
}

//nolint:contextcheck
func NewMaxActiveRPCLimiterStreamServerInterceptor(lr *limiting.Limiter) grpc.StreamServerInterceptor {
	return func(srv any, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
		if tag, ok := tagging.IOTagFromContext(ss.Context()); ok && tag == IOTagCritical.String() {
			return handler(srv, ss)
		}

		release, ok := lr.TryAcquire(info.FullMethod)
		if !ok {
			return new(apistatus.ResourceExhausted)
		}
		defer release()

		return handler(srv, ss)
	}
}