package buck

import (
	"context"
	"time"
)

// Bucket implements a token bucket rate limiter.
type Bucket struct {
	tokens chan struct{}
}

// Init initializes an existing Bucket with the given burst capacity.
func (b *Bucket) Init(burst int) {
	*b = Bucket{
		tokens: make(chan struct{}, burst),
	}

	for i := 0; i < burst; i++ {
		b.tokens <- struct{}{}
	}
}

// New is like Init but also allocates
func New(burst int) *Bucket {
	var b Bucket
	(&b).Init(burst)
	return &b
}

// Refill adds tokens to the bucket at the specified rate until the context is cancelled.
func (b Bucket) Refill(ctx context.Context, rate time.Duration) {
	ticker := time.NewTicker(rate)
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			select {
			case b.tokens <- struct{}{}:
			default:
				// Bucket full
			}
		}
	}
}

// Wait blocks until a token is available or the context is cancelled.
func (b Bucket) Wait(ctx context.Context) error {
	select {
	case <-b.tokens:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}