package waiterr

import (
	"errors"
	"sync"
)

// WaitErr wraps a sync.WaitGroup with error handling.
type WaitErr struct {
	wg   sync.WaitGroup
	errs []error
	mut  sync.RWMutex
}

// Go runs f in its own goroutine. When f returns, its error is stored, and returned
// with we.Wait()
func (we *WaitErr) Go(f func() error) {
	wrap := func() {
		err := f()
		we.mut.Lock()
		defer we.mut.Unlock()
		we.errs = append(we.errs, err)
	}
	we.wg.Go(wrap)
}

// WaitForError waits for the first error to be returned by one of our go routines, and immediately returns
// with that error. If all functions return successfully, a nil is returned.
func (we *WaitErr) WaitForError() error {
	// Implement this
	// If we already have an error, return it immediately without waiting
	// If no error has yet been returned, wait for the very first error and return it
	return nil
}

// Wait for all current goroutines to finish. Return an error that combines all errors returned
// in the group so far (if any).
func (we *WaitErr) Wait() error {
	we.wg.Wait()
	we.mut.RLock()
	defer we.mut.RUnlock()
	return errors.Join(we.errs...)
}

// Unwrap returns all non-nil errors returned by our functions.
// If we.errs is empty, or all errors are nil, just return nil.
func (we *WaitErr) Unwrap() []error {
	errs := make([]error, 0, len(we.errs))
	for _, e := range we.errs {
		if e != nil {
			errs = append(errs, e)
		}
	}
	if len(errs) == 0 {
		return nil
	}
	return errs
}