package main

import (
	"fmt"
	"sync"
	"time"
)

// Generator produces numbers
func generate(done <-chan struct{}, nums ...int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for _, n := range nums {
			select {
			case out <- n:
			case <-done:
				return
			}
		}
	}()
	return out
}

// Stage 1: Squarer
func square(done <-chan struct{}, in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			select {
			case out <- n * n:
			case <-done:
				return
			}
		}
	}()
	return out
}

// Stage 2: Multiplier
func multiplyBy(done <-chan struct{}, in <-chan int, factor int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for n := range in {
			select {
			case out <- n * factor:
			case <-done:
				return
			}
		}
	}()
	return out
}

func main() {
	done := make(chan struct{})
	defer close(done) // Signal all goroutines to stop when main exits

	// Set up the pipeline
	input := generate(done, 1, 2, 3, 4, 5)
	squared := square(done, input)
	result := multiplyBy(done, squared, 2)

	// Consume the final results
	fmt.Println("Pipeline Results:")
	for r := range result {
		fmt.Println(r)
	}
	// Expected: 2, 8, 18, 32, 50
	time.Sleep(10 * time.Millisecond) // Give goroutines a chance to clean up
}