package util

import (
	"time"
)

type RateLimit struct {
	list   *Queue
	td     time.Duration
	num    uint64
	opened bool
}

// NewRateLimit 获得一个实例，每 sec 秒 可以处理 num 次
func NewRateLimit(num uint64, td time.Duration) (rl *RateLimit) {
	if num == 0 {
		num = 10
	}
	lst := NewQueue(num)
	if td == 0 {
		td = time.Second
	}
	rl = &RateLimit{
		list: lst,
		num:  num,
		td:   td,
	}
	rl.Open()
	return
}

// SetDuration 动态改变间隔参数
func (rl *RateLimit) SetDuration(td time.Duration) {
	if td == 0 {
		td = time.Second
	}
	if td != rl.td {
		rl.td = td
	}
}

// Open 启动
func (rl *RateLimit) Open() {
	if rl.opened {
		return
	}
	rl.opened = true
	rl.list.Pause(false)
	go func() {
		for td := time.Duration(uint64(rl.td) / rl.num); rl.opened; {
			rl.list.Push(struct{}{})
			time.Sleep(td)
		}
	}()
}

// Close 关闭
func (rl *RateLimit) Close() {
	if !rl.opened {
		return
	}
	rl.list.Pause(true)
	rl.opened = false
}

// Status 当前可用量
func (rl *RateLimit) Status() int {
	return len(rl.list.data)
}

// Check 是否可用，非阻塞
func (rl *RateLimit) Check() bool {
	if !rl.opened {
		return true
	}
	_, ok := rl.list.Pop()
	return ok
}

// CheckWait 等待可用，阻塞，可设置超时
func (rl *RateLimit) CheckWait(td time.Duration) bool {
	if !rl.opened {
		return true
	}
	_, err := rl.list.PopWait(td)
	return err == nil
}