package klog

import "time"

type RateLimit struct {
	slotsLen           int
	slots              []int
	limitPerSec        int
	limitPerSlotPerSec float64
	previousI          int
	previousTime       int64
}

func NewRateLimit(limitPerSec, slotsPerSec int) *RateLimit {
	if limitPerSec < 1 {
		limitPerSec = 1
	}

	if slotsPerSec < 0 {
		slotsPerSec = 1
	} else if slotsPerSec > 10 {
		slotsPerSec = 10
	}

	return &RateLimit{
		slotsLen:           slotsPerSec,
		slots:              make([]int, slotsPerSec, slotsPerSec),
		limitPerSec:        limitPerSec,
		limitPerSlotPerSec: float64(limitPerSec) / float64(slotsPerSec),
		previousI:          0,
		previousTime:       time.Now().UnixMilli(),
	}
}

func (o *RateLimit) Wait() {
	added := false
	for {
		// 距离上次调用超过1秒，清零计数器
		now := time.Now().UnixMilli()
		if now > o.previousTime+1000 {
			for i := range o.slots {
				o.slots[i] = 0
			}
		}

		// 每隔一段时间切换一次slot
		i := int(now * int64(o.slotsLen) / 1000 % int64(o.slotsLen))
		if o.previousI != i {
			o.slots[i] = 0
			o.previousI = i
		}

		//计数
		if !added {
			o.slots[i] += 1
			added = true
		}

		// 每秒总数、每slot计数分别不得超限
		if sum(o.slots) >= o.limitPerSec || float64(o.slots[i]) > o.limitPerSlotPerSec && o.slots[i] > 0 {
			// 如果超限，休眠一个slot的跨度
			time.Sleep(time.Millisecond * time.Duration(1000/o.slotsLen))
		} else {
			break
		}

		o.previousTime = now
	}
}

func sum(slice []int) (total int) {
	for i := range slice {
		total += slice[i]
	}
	return total
}