package main

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

// job passed to goroutines. blend color from img1 and img2 at position (x, y)
type blendColorJob struct {
	X    int
	Y    int
	Img1 image.Image
	Img2 image.Image
}

// new color after blend, to apply at position (x, y)
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) BlendImagesMain() {
	log.Println("Blending the images...")

	p.InitWorkerPool()

	// output image, ready to receive pixel values
	dimensions := dimensionsToRectangle(ConfigRegister.OutputWidth, ConfigRegister.OutputHeight)
	outImg := image.NewRGBA(dimensions)

	for i := 1; i < ConfigRegister.Amount; i++ {
		// get two random images and load them
		var img1 image.Image
		if ConfigRegister.BaseImage != "" && i == 1 {
			img1 = loadImage(ConfigRegister.BaseImage)
		} else {
			if i != 1 {
				img1 = loadImage("output.jpg")
			} else {
				img1 = loadImage(getRandomImage())
			}
		}

		img2 := loadImage(getRandomImage())

		// pool of workers unionizing, ready to blend a new picture using the power of friendship
		var pool blendWorkerPool
		pool.InitWorkerPool()

		// main blending routine
		pool.BlendImages(img1, img2, outImg)

		encodeImage(outImg)
	}
}

func (p *blendWorkerPool) BlendImages(img1 image.Image, img2 image.Image, out *image.RGBA) {
	dimensions := dimensionsToRectangle(ConfigRegister.OutputWidth, ConfigRegister.OutputHeight)

	// TODO: use a worker pool for those operations ?
	// resize image
	img1Resized := resize(img1)
	img2Resized := resize(img2)

	p.RunWorkers(out)

	p.SendBlendJobs(dimensions, img1Resized, img2Resized)

	// 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()
}

// encode the image
func encodeImage(imgData *image.RGBA) {
	outputFile := fmt.Sprintf("%s/%s", ConfigRegister.OutputDir, "output.jpg")
	out, _ := os.Create(outputFile)
	defer out.Close()
	log.Println("Encoding the new image...")
	jpeg.Encode(out, imgData, nil)
	log.Println("Done.")
}

// convert RGBA pixel to grayscale. BT.709 luminosity formula
func grayscale(pixel color.Color) color.Color {
	c := color.RGBAModel.Convert(pixel).(color.RGBA)
	gray := uint8(0.2126*float64(c.R) + 0.7152*float64(c.G) + 0.0722*float64(c.B))
	return color.RGBA{
		R: gray,
		G: gray,
		B: gray,
		A: c.A,
	}
}

// sepia toning filter
func sepia(pixel color.Color) color.Color {
	c := color.RGBAModel.Convert(pixel).(color.RGBA)

	var npr uint8
	var npg uint8
	var npb uint8

	tr := float64(c.R)*0.393 + float64(c.G)*0.769 + float64(c.B)*0.189
	tg := float64(c.R)*0.349 + float64(c.G)*0.686 + float64(c.B)*0.168
	tb := float64(c.R)*0.272 + float64(c.G)*0.534 + float64(c.B)*0.131

	npr = uint8(math.Min(255, tr))
	npg = uint8(math.Min(255, tg))
	npb = uint8(math.Min(255, tb))

	return color.RGBA{
		R: npr,
		G: npg,
		B: npb,
		A: c.A,
	}
}

// if a filter is specified, apply it to the pixel, if not returns it unchanged
func filter(pixel color.Color) color.Color {
	if ConfigRegister.Grayscale {
		return grayscale(pixel)
	}

	if ConfigRegister.Sepia {
		return sepia(pixel)
	}
	return pixel
}

func average(fc uint8, bc uint8, fa uint8) float64 {
	return (float64(fc)/2 + float64(bc)/2)
}

func darken(fc uint8, bc uint8, fa uint8) float64 {
	return math.Min(float64(fc), float64(bc))
}

func lighten(fc uint8, bc uint8, fa uint8) float64 {
	return math.Max(float64(fc), float64(bc))
}

func multiply(fc uint8, bc uint8, fa uint8) float64 {
	return float64(fc) * float64(bc) / 255
}

// produce absolute garbage
func fuckyfun(fc uint8, bc uint8, fa uint8) float64 {
	return float64((fc * fa) + (bc * (fa * 2)))
}

// get the blending method from the config register
func blend(fc uint8, bc uint8, fa uint8, ba uint8) float64 {
	var newValue float64
	switch ConfigRegister.Method {
	case "darken":
		newValue = darken(fc, bc, fa)
	case "average":
		newValue = average(fc, bc, fa)
	case "lighten":
		newValue = lighten(fc, bc, fa)
	case "multiply":
		newValue = multiply(fc, bc, fa)
	case "fuckyfun":
		newValue = fuckyfun(fc, bc, fa)
	default:
		newValue = darken(fc, bc, fa)
	}
	return newValue
}

// 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
	new := color.RGBA{
		R: r, G: g, B: b, A: a,
	}
	return filter(new)
}

// creates a new rectangle with the min height and width from both images
func getMaxDimensions(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}
}

// resizes image using dimensions from config register (nearest neighbour interpolation)
func resize(img image.Image) (resized *image.RGBA) {
	imgSize := img.Bounds().Size()
	dimensions := dimensionsToRectangle(ConfigRegister.OutputWidth, ConfigRegister.OutputHeight)
	xscale := float64(imgSize.X) / float64(dimensions.Max.X)
	yscale := float64(imgSize.Y) / float64(dimensions.Max.Y)
	// creates new rescaled image based on a given image dimensions
	resized = image.NewRGBA(dimensions)
	// get pixels from the original image
	for x := 0; x < dimensions.Max.X; x++ {
		for y := 0; y < dimensions.Max.Y; y++ {
			xp := int(math.Floor(float64(x) * xscale))
			yp := int(math.Floor(float64(y) * yscale))
			pixel := img.At(xp, yp)
			resized.Set(x, y, pixel)
		}
	}
	return
}

// returns a Rectangle of dimensions <width>x<height>
func dimensionsToRectangle(width int, height int) image.Rectangle {
	upLeft := image.Point{0, 0}
	lowRight := image.Point{width, height}
	return image.Rectangle{upLeft, lowRight}
}