package main

import (
	"context"
	"errors"
	"fmt"
	"golang.org/x/sync/errgroup"
	"slices"
	"sync"
)

// Squares owns every goroutine it starts and returns only after all exit.
func Squares(parent context.Context, values []int, workers int) ([]int, error) {
	if workers <= 0 {
		return nil, errors.New("workers must be positive")
	}
	ctx, cancel := context.WithCancel(parent)
	defer cancel()
	jobs := make(chan int)
	results := make(chan int)
	var all, producers sync.WaitGroup
	all.Add(1)
	go func() {
		defer all.Done()
		defer close(jobs)
		for _, n := range values {
			select {
			case jobs <- n:
			case <-ctx.Done():
				return
			}
		}
	}()
	for range workers {
		producers.Add(1)
		all.Add(1)
		go func() {
			defer producers.Done()
			defer all.Done()
			for {
				select {
				case <-ctx.Done():
					return
				case n, ok := <-jobs:
					if !ok {
						return
					}
					select {
					case results <- n * n:
					case <-ctx.Done():
						return
					}
				}
			}
		}()
	}
	all.Add(1)
	go func() { defer all.Done(); producers.Wait(); close(results) }()
	var out []int
	for n := range results {
		out = append(out, n)
	}
	all.Wait()
	if err := parent.Err(); err != nil {
		return nil, err
	}
	slices.Sort(out)
	return out, nil
}

// The broker is single-owner: only its caller touches subscribers and closes them.
func Publish(subs []chan int, value int) (dropped int) {
	for _, sub := range subs {
		select {
		case sub <- value:
		default:
			dropped++
		}
	}
	return
}
func CancelPeers() error {
	g, ctx := errgroup.WithContext(context.Background())
	started := make(chan struct{})
	g.Go(func() error { close(started); <-ctx.Done(); return ctx.Err() })
	g.Go(func() error { <-started; return errors.New("stage failed") })
	return g.Wait()
}
func main() {
	out, err := Squares(context.Background(), []int{1, 2, 3, 4}, 2)
	fmt.Println(out, err)
	fast, slow := make(chan int, 1), make(chan int, 1)
	subs := []chan int{fast, slow}
	fmt.Println("first dropped:", Publish(subs, 1))
	fmt.Println("fast:", <-fast)
	fmt.Println("second dropped:", Publish(subs, 2))
	close(fast)
	close(slow)
	fmt.Println("remaining:", <-fast, <-slow)
	fmt.Println("errgroup:", CancelPeers())
}