package main

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

// Worker represents a single worker goroutine
type Worker struct {
	ID int
}

func (w *Worker) Start(tasks <-chan int, ready chan<- chan int, results chan<- string, wg *sync.WaitGroup) {
	defer wg.Done()
	fmt.Printf("Worker %d started
", w.ID)

	myTaskCh := make(chan int) // Worker's specific input channel

	for {
		select {
		case ready <- myTaskCh: // Signal readiness by sending my task channel to the dispatcher
			task, ok := <-myTaskCh // Wait for a task on my specific channel
			if !ok { // myTaskCh closed by dispatcher
				fmt.Printf("Worker %d shutting down (task channel closed)
", w.ID)
				return
			}
			fmt.Printf("Worker %d processing task %d
", w.ID, task)
			time.Sleep(time.Duration(task % 3) * 50 * time.Millisecond) // Simulate variable work
			results <- fmt.Sprintf("Worker %d completed task %d", w.ID, task)

		case <-time.After(5 * time.Second): // Optional: timeout if worker is idle for too long
			fmt.Printf("Worker %d idle for 5s, shutting down.
", w.ID)
			return
		}
	}
}

// Dispatcher manages distributing tasks to available workers
func Dispatcher(numWorkers int, tasks <-chan int, ready <-chan chan int, results chan<- string, wg *sync.WaitGroup) {
	// Start workers
	for i := 0; i < numWorkers; i++ {
		w := &Worker{ID: i}
		wg.Add(1)
		go w.Start(tasks, ready, results, wg)
	}

	// Distribute tasks as workers become ready
	activeWorkerTaskChannels := make([]chan int, 0, numWorkers)
	for task := range tasks {
		workerTaskCh := <-ready // Get an available worker's task channel
		activeWorkerTaskChannels = append(activeWorkerTaskChannels, workerTaskCh)
		workerTaskCh <- task    // Send task to that specific worker
	}

	// After all tasks are sent, close all active worker task channels to signal them to shut down
	for _, ch := range activeWorkerTaskChannels {
		close(ch)
	}
	fmt.Println("Dispatcher finished sending tasks and closed worker channels.")
	close(results) // Close results channel as well, after all workers are signalled to exit.
}

func main() {
	const numWorkers = 3
	const numTasks = 10

	tasksCh := make(chan int, numTasks) // Channel for incoming tasks
	readyCh := make(chan chan int)       // Channel for workers to signal readiness
	resultsCh := make(chan string)       // Unbuffered for demonstration, or buffered for higher throughput

	var workerWg sync.WaitGroup // To wait for individual workers

	// Start the dispatcher (which in turn starts workers)
	go Dispatcher(numWorkers, tasksCh, readyCh, resultsCh, &workerWg)

	// Simulate task production
	for i := 0; i < numTasks; i++ {
		tasksCh <- i
		time.Sleep(20 * time.Millisecond)
	}
	close(tasksCh) // No more tasks to send

	// Collect results
	for result := range resultsCh {
		fmt.Println("Main received result:", result)
	}

	workerWg.Wait() // Wait for all workers to finish their tasks and exit

	fmt.Println("All tasks processed and workers shut down.")
}