// Exercise 02 — Worker Pool
//
// ASSIGNMENT
// ----------
// Implement Pool so that it:
//
//  1. Spawns exactly `workers` goroutines.
//  2. Each worker consumes jobs from `jobs` until the channel is closed.
//  3. Results of fn(job) are sent to the returned channel.
//  4. The returned channel is CLOSED after ALL workers finish.
//     (Whoever ranges over the output must terminate — no deadlocks.)
//  5. No goroutines may leak.
//
// Run the tests:   go test -race ./02-worker-pool/
package main

import (
	"fmt"
	"sync"
)

// Pool fans `jobs` out to `workers` goroutines running fn, fans results in.
// TODO: implement me.
func Pool[T, R any](jobs <-chan T, workers int, fn func(T) R) <-chan R {
	out := make(chan R)
	var wg sync.WaitGroup
	// start all the workers and increment the waitgroup
	for i := 0; i < workers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for j := range jobs {
				out <- fn(j)
			}
		}()
	}
	// TODO:
	//   - start `workers` goroutines, each: for j := range jobs { out <- fn(j) }
	//   - close(out) exactly once, only after every worker returned
	//     (hint: sync.WaitGroup + one extra "closer" goroutine)

	// start go routine that closes out channel when all workers are done
	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func main() {
	jobs := make(chan int)
	go func() {
		for i := 1; i <= 9; i++ {
			jobs <- i
		}
		close(jobs)
	}()
	for r := range Pool(jobs, 3, func(n int) int { return n * n }) {
		fmt.Println(r)
	}
}