package ratecache

import (
	"sync"
	"time"

	"golang.org/x/time/rate"
)

const defaultBurst = 30

var defaultLimit = rate.Every(time.Second)

// Store holds rate limiters (token buckets) by key.
type Store struct {
	// Limit is how often events are allowed. Defaults to one per second.
	Limit rate.Limit
	// Burst is the maximum number of allowed tokens that can be consumed. It is the initial size
	// of the token bucket. Defaults to 30
	Burst int

	limiters map[string]*rate.Limiter
	mu       sync.Mutex
}

// CreateOrGet lazy-creates a rate limiter and returns it. The rate limiter is stored
// for future lookups.
func (s *Store) CreateOrGet(key string) *rate.Limiter {
	s.mu.Lock()
	defer s.mu.Unlock()

	if s.limiters == nil {
		s.limiters = make(map[string]*rate.Limiter)
	}

	if limiter, ok := s.limiters[key]; ok {
		return limiter
	}

	limiter := rate.NewLimiter(s.limit(), s.burst())
	s.limiters[key] = limiter

	return limiter
}

func (s *Store) limit() rate.Limit {
	if s.Limit <= 0 {
		return defaultLimit
	}

	return s.Limit
}

func (s *Store) burst() int {
	if s.Burst <= 0 {
		return defaultBurst
	}

	return s.Burst
}
