package ac

import (
	"context"
	"jiby.grabit/src/ac/buck"
	"time"
)

type sleeper interface {
	Wait(context.Context) error
}

type Worker[X any] struct {
	Input  <-chan X
	Output chan<- X
	s      sleeper
}

// InitBucket initializes the worker with an access control bucket.
func (w *Worker[X]) InitBucket(rate time.Duration, burst int) {
	var bucket buck.Bucket
	bucket.Init(burst)
	go bucket.Refill(context.TODO(), rate)
	w.s = bucket
}

func (w *Worker[X]) Work(ctx context.Context) error {
	select {
	case <-ctx.Done():
		return ctx.Err()

	case val, ok := <-w.Input:
		if !ok {
			return nil
		}

		if err := w.s.Wait(ctx); err != nil {
			return err
		}

		select {
		case w.Output <- val:
			return nil
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}