package main

import (
	"fmt"
	"image"
	"image/color"
	"image/png"
	"math/rand"
	"os"
	"runtime"
	"sync"
	"time"
)

// Task represents a generic task to be executed
type Task func() Result

// Result represents the result of processing a task
type Result struct {
	filename string
	duration time.Duration
}

// generateImage creates a random image
func generateImage(filename string, width, height int) Task {
	return func() Result {
		start := time.Now()

		img := image.NewRGBA(image.Rect(0, 0, width, height))
		for y := 0; y < height; y++ {
			for x := 0; x < width; x++ {
				img.Set(x, y, color.RGBA{
					R: uint8(rand.Intn(256)),
					G: uint8(rand.Intn(256)),
					B: uint8(rand.Intn(256)),
					A: 255,
				})
			}
		}
		f, err := os.Create(filename)
		if err != nil {
			fmt.Printf("Error creating file %s: %v\n", filename, err)
			return Result{filename, 0}
		}
		defer f.Close()

		err = png.Encode(f, img)
		if err != nil {
			fmt.Printf("Error encoding image %s: %v\n", filename, err)
			return Result{filename, 0}
		}

		return Result{filename, time.Since(start)}
	}
}

// processImage applies a simple filter to the image
func processImage(filename string) Task {
	return func() Result {
		start := time.Now()

		// Open the image file
		file, err := os.Open(filename)
		if err != nil {
			fmt.Printf("Error opening file %s: %v\n", filename, err)
			return Result{filename, 0}
		}
		defer file.Close()

		// Decode the image
		img, err := png.Decode(file)
		if err != nil {
			fmt.Printf("Error decoding image %s: %v\n", filename, err)
			return Result{filename, 0}
		}

		// Apply a simple filter (invert colors)
		bounds := img.Bounds()
		inverted := image.NewRGBA(bounds)
		for y := bounds.Min.Y; y < bounds.Max.Y; y++ {
			for x := bounds.Min.X; x < bounds.Max.X; x++ {
				oldColor := img.At(x, y)
				r, g, b, a := oldColor.RGBA()
				inverted.Set(x, y, color.RGBA{
					R: uint8(255 - r>>8),
					G: uint8(255 - g>>8),
					B: uint8(255 - b>>8),
					A: uint8(a >> 8),
				})
			}
		}

		// Save the processed image
		outFile, err := os.Create("processed_" + filename)
		if err != nil {
			fmt.Printf("Error creating output file for %s: %v\n", filename, err)
			return Result{filename, 0}
		}
		defer outFile.Close()

		err = png.Encode(outFile, inverted)
		if err != nil {
			fmt.Printf("Error encoding processed image %s: %v\n", filename, err)
			return Result{filename, 0}
		}

		return Result{filename, time.Since(start)}
	}
}

func worker(id int, tasks <-chan Task, results chan<- Result, wg *sync.WaitGroup) {
	defer wg.Done()
	for task := range tasks {
		fmt.Printf("Worker %d processing new task...\n", id)
		results <- task()
	}
}

const MAX_WORKERS = 4

func executeTasks(tasks []Task) time.Duration {
	numWorkers := runtime.NumCPU()
	if numWorkers > MAX_WORKERS {
		numWorkers = MAX_WORKERS
	}
	tasksChan := make(chan Task, len(tasks))
	resultsChan := make(chan Result, len(tasks))

	// Fill task channel
	for _, task := range tasks {
		tasksChan <- task
	}
	close(tasksChan)

	// Start worker pool
	var wg sync.WaitGroup
	for i := 0; i < numWorkers; i++ {
		wg.Add(1)
		go worker(i, tasksChan, resultsChan, &wg)
	}

	// Wait for all workers to finish
	go func() {
		wg.Wait()
		close(resultsChan)
	}()

	// Collect results
	var totalDuration time.Duration
	for result := range resultsChan {
		fmt.Printf("Processed %s in %v\n", result.filename, result.duration)
		totalDuration += result.duration
	}

	return totalDuration
}

func main() {
	numImages := 20
	var generateTasks []Task
	var processTasks []Task

	// Create image generation tasks
	for i := 0; i < numImages; i++ {
		filename := fmt.Sprintf("image_%d.png", i)
		generateTasks = append(generateTasks, generateImage(filename, 1000, 1000))
	}

	// Execute image generation tasks
	fmt.Println("Generating images...")
	genDuration := executeTasks(generateTasks)
	fmt.Printf("All images generated in %v\n", genDuration)
	fmt.Printf("Average time per image generation: %v\n", genDuration/time.Duration(numImages))

	// Create image processing tasks
	for i := 0; i < numImages; i++ {
		filename := fmt.Sprintf("image_%d.png", i)
		processTasks = append(processTasks, processImage(filename))
	}

	// Execute image processing tasks
	fmt.Println("\nProcessing images...")
	procDuration := executeTasks(processTasks)
	fmt.Printf("All images processed in %v\n", procDuration)
	fmt.Printf("Average time per image processing: %v\n", procDuration/time.Duration(numImages))

	fmt.Printf("\nTotal time: %v\n", genDuration+procDuration)
}