package ratelimit

import (
	"context"
	"sync"
	"time"

	"github.com/juju/ratelimit"
	"github.com/tidbcloud/pkgv2/bizerr"
	"github.com/tidbcloud/serverless-service/internal/utils"
	"google.golang.org/grpc"
)

type RateLimiter interface {
	Limit(ctx context.Context) error
}

type rateLimiter struct {
	fillInterval time.Duration
	capacity     int64
	limiters     *sync.Map
}

func (r *rateLimiter) Limit(ctx context.Context) error {
	id := utils.GetIdentityKeyFromContext(ctx)
	if v, ok := r.limiters.Load(id); ok {
		if v.(*ratelimit.Bucket).TakeAvailable(1) == 0 {
			return bizerr.New(bizerr.ErrCodeTooManyRequests, "Too many requests")
		}
	} else {
		limiter := ratelimit.NewBucketWithQuantum(r.fillInterval, r.capacity, r.capacity)
		limiter.TakeAvailable(1)
		r.limiters.Store(id, limiter)
	}

	return nil
}

func NewRateLimiter(interval time.Duration, capacity int64) RateLimiter {
	return &rateLimiter{
		fillInterval: interval,
		capacity:     capacity,
		limiters:     &sync.Map{},
	}
}

// UnaryServerInterceptor returns a new unary server interceptor that performs rate limiting
func UnaryServerInterceptor(r RateLimiter) grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
		if err := r.Limit(ctx); err != nil {
			return nil, err
		}
		return handler(ctx, req)
	}
}