package ratelimit

import (
	"context"
	"gitee.com/mosache/YFrame/micro/ratelimit/internal/errs"
	"google.golang.org/grpc"
	"sync/atomic"
	"time"
)

type FixWindowLimiter struct {
	timeUnix     int64 // 窗口开始时间
	timeDuration int64 // 时间窗口长度
	cnt          int64 // 请求计数
	rate         int64 // 在这个窗口内，允许通过的最大请求数量
}

func NewFixWindowLimiter(timeDuration time.Duration, rate int64) *FixWindowLimiter {
	return &FixWindowLimiter{
		timeUnix:     time.Now().UnixNano(),
		timeDuration: timeDuration.Nanoseconds(),
		cnt:          0,
		rate:         rate,
	}
}

func (f *FixWindowLimiter) Build() grpc.UnaryServerInterceptor {
	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) {
		now := time.Now().UnixNano()
		timeUnix := atomic.LoadInt64(&f.timeUnix)
		cnt := atomic.LoadInt64(&f.cnt)
		if now > timeUnix+f.timeDuration { // 重置
			if atomic.CompareAndSwapInt64(&f.timeUnix, timeUnix, now) {
				atomic.CompareAndSwapInt64(&f.cnt, cnt, 0)
			}
		}

		cnt = atomic.AddInt64(&f.cnt, 1)

		if cnt > f.rate {
			err = errs.ErrLimiterOverCapacity
			return
		}
		resp, err = handler(ctx, req)
		return
	}
}