package errgroup

import (
	"context"
	"sync"
)

type Group struct {
	ctx    context.Context
	cancel context.CancelFunc

	wg sync.WaitGroup

	firstErr chan error

	// Used to limit number of goroutines (workers) for group.
	workerLimCh chan struct{}
}

func WithContext(parent context.Context, workers int) (*Group, context.Context) {
	ctx, cancel := context.WithCancel(parent)
	return &Group{
		ctx:         ctx,
		cancel:      cancel,
		firstErr:    make(chan error, 1),
		workerLimCh: make(chan struct{}, workers),
	}, ctx
}

func (g *Group) Go(f func(ctx context.Context) error) {
	// Wait for worker to be freed.
	g.workerLimCh <- struct{}{}

	g.wg.Add(1)
	go func() {
		defer func() {
			g.wg.Done()

			<-g.workerLimCh // release worker
		}()

		if err := f(g.ctx); err != nil {
			select {
			case g.firstErr <- err:
				g.cancel()
			default: // got error from other goroutine
			}
		}
	}()
}

func (g *Group) Wait() error {
	g.wg.Wait()
	g.cancel()

	select {
	case err := <-g.firstErr:
		return err
	default:
		return nil
	}
}