package main

import (
	"fmt"
	"image"
	"image/color"
	"image/jpeg"
	"math"
	"os"
	"runtime"
	"sync"
)

type blendColorJob struct {
	X    int
	Y    int
	Img1 image.Image
	Img2 image.Image
}

type blendColorResult struct {
	X     int
	Y     int
	Color color.Color
}

type blendWorkerPool struct {
	CpuCores   int
	Jobs       chan blendColorJob
	Results    chan blendColorResult
	WorkerWG   *sync.WaitGroup
	ConsumerWG *sync.WaitGroup
}

func (p *blendWorkerPool) InitWorkerPool() {
	p.CpuCores = runtime.NumCPU()
	p.Jobs = make(chan blendColorJob, 1)
	p.Results = make(chan blendColorResult, 1)
	p.WorkerWG = new(sync.WaitGroup)
	p.ConsumerWG = new(sync.WaitGroup)
}

// runs the pool of goroutines
func (p *blendWorkerPool) RunWorkers(outImg *image.RGBA) {
	for i := 1; i < p.CpuCores; i++ {
		p.WorkerWG.Add(1)
		p.ConsumerWG.Add(1)
		go p.BlendWorker()
		go p.SetImageWorker(outImg)
	}
}

// blending pixels and assigning them to the output image
func (p *blendWorkerPool) SendBlendJobs(dimensions image.Rectangle, img1 image.Image, img2 image.Image) {
	for x := 0; x < dimensions.Max.X; x++ {
		for y := 0; y < dimensions.Max.Y; y++ {
			p.Jobs <- blendColorJob{x, y, img1, img2}
		}
	}
	close(p.Jobs)
}

// goroutine to blend colors at position (x, y) from two images
func (p *blendWorkerPool) BlendWorker() {
	defer p.WorkerWG.Done()
	for j := range p.Jobs {
		c1 := j.Img1.At(j.X, j.Y)
		c2 := j.Img2.At(j.X, j.Y)
		c3 := blendColor(c1, c2)
		p.Results <- blendColorResult{j.X, j.Y, c3}
	}
}

// goroutine to set the new color values at newImage position (x, y)
func (p *blendWorkerPool) SetImageWorker(newImage *image.RGBA) {
	defer p.ConsumerWG.Done()
	for r := range p.Results {
		newImage.Set(r.X, r.Y, r.Color)
	}
}

func (p *blendWorkerPool) BlendImages(img1 image.Image, img2 image.Image) {
	dimensions := getDimensions(img1, img2)
	// output image, ready to receive pixel values
	outImg := image.NewRGBA(dimensions)
	p.RunWorkers(outImg)

	p.SendBlendJobs(dimensions, img1, img2)

	// first waitgroup to wait for the results to be ready before closing the channel
	p.WorkerWG.Wait()
	close(p.Results)

	// second waitgroup to ensure the encoding doesn't start before the goroutines are done
	p.ConsumerWG.Wait()
	encodeImage(outImg)
}

// encode the image
func encodeImage(imgData *image.RGBA) {
	out, _ := os.Create("output.jpg")
	defer out.Close()
	fmt.Print("Encoding the image...")
	jpeg.Encode(out, imgData, nil)
	fmt.Println(" Done.")
}

// convert RGBA pixel to grayscale
func grayscale(pixel color.Color) color.Color {
	c := color.RGBAModel.Convert(pixel).(color.RGBA)
	gray := uint8((c.R + c.G + c.B) / 3)
	return color.RGBA{
		R: gray,
		G: gray,
		B: gray,
		A: c.A,
	}
}

// alpha blending RGBA pixels
func blend(fc uint8, bc uint8, fa uint8, ba uint8) float64 {
	// fc : foreground color
	// bc : background color
	// fa : foreground alpha
	// ba : background alpha
	// all values are alpha-premultiplied

	// darken only
	return math.Min(float64(fc), float64(bc))

	// fucky fucky fun
	// return float64((fc * fa) + (bc * (fa * 2)))
}

// blend two RGBA colors/pixels and returns a new one
func blendColor(color1 color.Color, color2 color.Color) color.Color {
	oc1 := color.RGBAModel.Convert(color1).(color.RGBA)
	oc2 := color.RGBAModel.Convert(color2).(color.RGBA)
	r := uint8(blend(oc1.R, oc2.R, oc1.A, oc2.A))
	g := uint8(blend(oc1.G, oc2.G, oc1.A, oc2.A))
	b := uint8(blend(oc1.B, oc2.B, oc1.A, oc2.A))
	a := oc1.A + (1-oc1.A)*oc2.A
	return color.RGBA{
		R: r, G: g, B: b, A: a,
	}
}

// creates a new rectangle with the min height and width from both images
func getDimensions(img1 image.Image, img2 image.Image) image.Rectangle {
	// get dimensions for both images
	size1 := img1.Bounds().Size()
	size2 := img2.Bounds().Size()

	// final image sized from lowest width and height
	width := int(math.Min(float64(size1.X), float64(size2.X)))
	height := int(math.Min(float64(size1.Y), float64(size2.Y)))

	// the dimensions, as Points, of the output image
	upLeft := image.Point{0, 0}
	lowRight := image.Point{width, height}

	return image.Rectangle{upLeft, lowRight}
}