package ratelimit

import (
	"context"
	"github.com/zeromicro/go-zero/core/syncx"
	"google.golang.org/grpc"
	"time"
)

// 令牌桶
type TokenBucketLimiter struct {
	tokens chan struct{}
	ch     chan struct{}
}

func NewTokenBucketLimiter(capacity int, interval time.Duration) *TokenBucketLimiter {
	tokens := make(chan struct{}, capacity)
	closeCh := make(chan struct{})
	ticker := time.NewTicker(interval)
	go func() {
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				select {
				case tokens <- struct{}{}:
				default:
				}
			case <-closeCh:
				break
			default:
			}
		}
	}()
	return &TokenBucketLimiter{
		tokens: tokens,
		ch:     closeCh,
	}
}

func (t *TokenBucketLimiter) Build() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		select {
		case <-t.ch:
			/// 关闭了limiter，直接通过
			resp, err = handler(ctx, req)
		case <-t.tokens:
			resp, err = handler(ctx, req)
		case <-ctx.Done():
			err = ctx.Err()
		}

		return
	}
}

func (t *TokenBucketLimiter) Close() {
	syncx.Once(func() {
		close(t.ch)
	})
}