package main

import (
	"fmt"
	"image"
	"image/color"
	"image/jpeg"
	_ "image/jpeg"
	_ "image/png"
	"log"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"sync"
	"time"
)

// goroutine to blend colors at position (x, y) from two images
func blendWorker(jobs <-chan blendColorJob, results chan<- blendColorResult, wg *sync.WaitGroup) {
	defer wg.Done()
	for j := range jobs {
		c1 := j.Img1.At(j.X, j.Y)
		c2 := j.Img2.At(j.X, j.Y)
		c3 := blendColor(c1, c2)
		results <- blendColorResult{j.X, j.Y, c3}
	}
}

// goroutine to set the new color values at newImage position (x, y)
func setImageWorker(results <-chan blendColorResult, newImage *image.RGBA, wg *sync.WaitGroup) {
	defer wg.Done()
	for r := range results {
		newImage.Set(r.X, r.Y, r.Color)
	}
}

// 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,
	}
}

// load an image from a file into an Image value
func loadImage(filename string) image.Image {
	img, err := os.Open(filename)
	if err != nil {
		log.Fatal(err)
	}
	defer img.Close()

	imgData, _, err := image.Decode(img)
	if err != nil {
		log.Fatal(err)
	}

	return imgData
}

// Walk through a folder recursively and returns a list of image paths
func getImagesList(path string) []string {
	var imgs []string

	err := filepath.Walk(path,
		func(path string, info os.FileInfo, err error) error {
			ext := strings.ToLower(filepath.Ext(path))
			if err != nil {
				return err
			}
			if ext == ".jpg" || ext == ".png" {
				imgs = append(imgs, path)
			}
			return nil
		})

	if err != nil {
		log.Println(err)
	}

	return imgs
}

// Randomly choose x number of image from a given folder
func getRandomImages(number int) []string {
	rand.Seed(time.Now().UnixNano())
	var images []string
	dir := getImagesList("/home/gator/Photos/")
	for i := 0; i < number; i++ {
		index := rand.Intn(len(dir))
		images = append(images, dir[index])
	}
	return images
}

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

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

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

func main() {
	cpu := runtime.NumCPU()
	// channels
	jobs := make(chan blendColorJob, 1)
	results := make(chan blendColorResult, 1)
	// waitgroups
	wgWorker := new(sync.WaitGroup)
	wgConsumer := new(sync.WaitGroup)

	// get two random images and load them
	imgs := getRandomImages(2)
	// imData1 := loadImage("assets/moutons.jpg")
	// imData2 := loadImage("assets/lavande.jpg")
	imData1 := loadImage(imgs[0])
	imData2 := loadImage(imgs[1])

	dimensions := getDimensions(imData1, imData2)

	// output image, ready to receive pixel values
	outImg := image.NewRGBA(dimensions)
	out, _ := os.Create("output.jpg")
	defer out.Close()

	for i := 1; i < cpu; i++ {
		wgConsumer.Add(1)
		wgWorker.Add(1)
		go blendWorker(jobs, results, wgWorker)
		go setImageWorker(results, outImg, wgConsumer)
	}

	// blending pixels and assigning them to the output image
	for x := 0; x < dimensions.Max.X; x++ {
		for y := 0; y < dimensions.Max.Y; y++ {
			jobs <- blendColorJob{x, y, imData1, imData2}
		}
	}
	close(jobs)

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

	// second waitgroup to ensure the encoding doesn't start before the goroutines are done
	wgConsumer.Wait()
	fmt.Print("Encoding the image...")
	jpeg.Encode(out, outImg, nil)
	fmt.Println(" Done.")
}