// Package rate_limit ...
//
// Description : rate_limit ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 2024-06-20 20:40
package rate_limit

import (
	"context"
	"git.zhangdeman.cn/zhangdeman/rate_limit/define"
	"golang.org/x/time/rate"
	"sync"
	"time"
)

var (
	MemoryClient = &Memory{
		lock:         &sync.RWMutex{},
		limiterTable: make(map[string]*rate.Limiter),
	}
)

// Memory ...
//
// Author : go_developer@163.com<白茶清欢>
//
// Date : 20:40 2024/6/20
type Memory struct {
	lock         *sync.RWMutex
	limiterTable map[string]*rate.Limiter
}

func (m *Memory) AllowN(ctx context.Context, limitCfg *define.LimitConfig, tokenCnt int) (bool, error) {
	limiter := m.getLimiter(limitCfg)
	return limiter.AllowN(time.Now(), tokenCnt), nil
}

func (m *Memory) Reset(ctx context.Context, limitCfg *define.LimitConfig) error {
	m.lock.Lock()
	defer m.lock.Unlock()
	delete(m.limiterTable, limitCfg.Key)
	return nil
}

func (m *Memory) getLimiter(limitCfg *define.LimitConfig) *rate.Limiter {
	m.lock.Lock()
	defer m.lock.Unlock()
	if _, exist := m.limiterTable[limitCfg.Key]; !exist {
		m.limiterTable[limitCfg.Key] = rate.NewLimiter(rate.Limit(limitCfg.Rate), limitCfg.Total)
	}
	return m.limiterTable[limitCfg.Key]
}