package main

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

// producer generates numbers and sends them to its output channel.
func producer(id int, out chan<- int, wg *sync.WaitGroup) {
	defer wg.Done()
	for i := 0; i < 3; i++ {
		value := id*100 + i
		fmt.Printf("Producer %d: Sending %d
", id, value)
		out <- value
		time.Sleep(time.Duration(id+1) * 50 * time.Millisecond) // Simulate varying work
	}
}

// merge takes multiple input channels and merges their outputs into a single output channel.
func merge(done <-chan struct{}, inputs ...<-chan int) <-chan int {
	var wg sync.WaitGroup
	merged := make(chan int)

	// Function to read from one input channel and send to the merged channel
	output := func(c <-chan int) {
		defer wg.Done()
		for n := range c {
			select {
			case merged <- n:
			case <-done:
				return // Exit if done signal received
			}
		}
	}

	wg.Add(len(inputs))
	for _, c := range inputs {
		go output(c)
	}

	// Start a goroutine to close the merged channel once all inputs are closed.
	go func() {
		wg.Wait()
		close(merged)
	}()

	return merged
}

func main() {
	const numProducers = 3
	inputChannels := make([]<-chan int, numProducers)
	var producerWg sync.WaitGroup
	done := make(chan struct{}) // Channel to signal graceful shutdown

	fmt.Println("Main: Starting producers...")
	for i := 0; i < numProducers; i++ {
		ch := make(chan int)
		inputChannels[i] = ch
		producerWg.Add(1)
		go producer(i+1, ch, &producerWg)
	}

	// Close producer channels once all producers are done
	go func() {
		producerWg.Wait()
		for _, ch := range inputChannels {
			close(ch.(chan int)) // Type assertion needed to close send-only channel
		}
		fmt.Println("Main: All producers finished and their channels closed.")
	}()

	// Merge the output of all producers into a single channel
	mergedStream := merge(done, inputChannels...)

	fmt.Println("Main: Reading from merged stream...")
	for val := range mergedStream {
		fmt.Printf("Main: Received from merged stream: %d
", val)
	}

	// In a real application, you might close `done` channel
	// after some timeout or external signal.
	// close(done) // Uncomment to trigger shutdown if merging is endless.

	fmt.Println("Main: All data from merged stream processed.")
}