package main

import (
	"fmt"
	"image/color"
	"image/png"
	"path/filepath"

	"github.com/disintegration/imaging"
	"github.com/urfave/cli/v2"
	"golang.org/x/image/tiff"
	"image"
	"image/jpeg"
	"io"
	"mime"
	"os"
	"path"
	"strings"
)

func main() {
	commands := []*cli.Command{
		// add commands here
	}

	app := &cli.App{
		Name:        "imgStack",
		Usage:       "A command line tool to denoise or give super resolution to an image by stacking multiple images together.",
		Version:     "v1.0.0",
		Description: "imgStack is a command line tool to denoise or give super resolution to an image by stacking multiple images together.",
		Commands:    commands,
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:  "output",
				Value: "./output.png",
				Usage: "path to output file (should be a png)",
				Aliases: []string{
					"o",
				},
			},
			&cli.BoolFlag{
				Name:    "super-resolution",
				Aliases: []string{"s", "sr"},
				Usage:   "Whether to use super resolution or not. If this is set to true, the image will be doubled in size.",
				Value:   false,
			},
		},
		Authors: []*cli.Author{
			{
				Name:  "Mason Payne",
				Email: "mason@masonitestudios.com",
			},
		},
		Copyright:              "2024 Masonite Studios LLC",
		UseShortOptionHandling: true,
		Action: func(c *cli.Context) error {
			var needsCleanup bool
			var cleanUpFiles []string
			var sourceFiles []string
			// use the first argument as the file name
			// we need to collect the list of discrete images
			fmt.Println("Args: ", c.Args().Slice())

			fileList, err := getListOfFiles(c.Args().Slice())
			if err != nil {
				return fmt.Errorf("error getting list of files | %w", err)
			}

			fmt.Println("Files to be used: ", fileList)

			// check if superResolution has been requested
			superResolution := c.Bool("super-resolution")

			if len(fileList) == 0 {
				return fmt.Errorf("no files provided")
			}

			// loop over each file
			for _, fileLocation := range fileList {
				// make sure the image is either jpeg, png or tiff
				fileMimeType := mime.TypeByExtension(path.Ext(fileLocation))
				if fileMimeType != "image/jpeg" && fileMimeType != "image/png" && fileMimeType != "image/tiff" {
					return fmt.Errorf("unsupported file type | %v", fileMimeType)
				}

				if superResolution {
					// resize the images
					tempLocation, err := resizeImage(fileLocation, 2)
					if err != nil {
						return fmt.Errorf("error resizing image | %w", err)
					}

					if tempLocation != fileLocation {
						fmt.Printf("Resized image to %s\n", tempLocation)
						needsCleanup = true
						fileLocation = tempLocation
					}
					sourceFiles = append(sourceFiles, fileLocation) // use the resized image

					fmt.Println("Needs cleanup: ", needsCleanup)
					if needsCleanup {
						cleanUpFiles = append(cleanUpFiles, tempLocation)
					}
				} else {
					sourceFiles = append(sourceFiles, fileLocation) // use the original image
				}
			}

			err = stackImages(sourceFiles, c.String("output"))
			if err != nil {
				return fmt.Errorf("error stacking images | %w", err)
			}

			if needsCleanup {
				for _, file := range cleanUpFiles {
					// remove the temp file
					err = cleanUpTempFile(file)
					if err != nil {
						return fmt.Errorf("error removing temp file | %w", err)
					}
				}
			}

			return nil
		},
	}

	err := app.Run(os.Args)
	if err != nil {
		fmt.Println(fmt.Errorf("error running app | %w", err))
		return
	}
}

func cleanUpTempFile(fileLocation string) error {
	// remove the temp file
	err := os.Remove(fileLocation)
	if err != nil {
		return fmt.Errorf("error removing temp file | %w", err)
	}
	return nil
}

func resizeImage(fileLocation string, multiple int) (string, error) {
	// load the image
	var img image.Image
	img, err := imaging.Open(fileLocation, imaging.AutoOrientation(true))
	if err != nil {
		return "", fmt.Errorf("error opening image | %w", err)
	}

	// double the size of the image
	rect := img.Bounds()
	width := rect.Max.X - rect.Min.X
	height := rect.Max.Y - rect.Min.Y
	x1 := multiple * width
	y1 := multiple * height

	croppedImage := imaging.Fill(img, x1, y1, imaging.Center, imaging.Linear) // technically not cropped because it will still have the same aspect ratio

	// save the image to a temp file
	tempFile, err := os.CreateTemp("", "imgStack*"+path.Ext(fileLocation))
	if err != nil {
		return "", fmt.Errorf("error creating temp file | %w", err)
	}
	defer tempFile.Close()

	// encode the image to the temp file
	fileMimeType := mime.TypeByExtension(path.Ext(fileLocation))
	if fileMimeType == "image/jpeg" {
		err = jpeg.Encode(tempFile, croppedImage, nil)
		if err != nil {
			return "", fmt.Errorf("error encoding resized jpeg image | %w", err)
		}
	}

	if fileMimeType == "image/png" {
		err = png.Encode(tempFile, croppedImage)
		if err != nil {
			return "", fmt.Errorf("error encoding resized png image | %w", err)
		}
	}

	if fileMimeType == "image/tiff" {
		err = tiff.Encode(tempFile, croppedImage, nil)
		if err != nil {
			return "", fmt.Errorf("error encoding resized tiff image | %w", err)
		}
	}

	// return the temp file location
	return tempFile.Name(), nil
}

// stackImages stacks the images using the average value of each pixel in sourceFiles and saves the result to output
func stackImages(sourceFiles []string, output string) error {
	// load the images
	var images []image.Image
	for _, fileLocation := range sourceFiles {
		// load the image
		var img image.Image
		img, err := imaging.Open(fileLocation, imaging.AutoOrientation(true))
		if err != nil {
			return fmt.Errorf("error opening image | %w", err)
		}
		images = append(images, img)
	}

	// get the bounds of the first image
	bounds := images[0].Bounds()

	// create a new image
	newImage := image.NewRGBA(bounds)

	// loop over each pixel in the new image
	for x := 0; x < bounds.Max.X; x++ {
		for y := 0; y < bounds.Max.Y; y++ {
			// get the average value of each pixel
			var r, g, b, a uint32
			for _, img := range images {
				// get the pixel at x, y
				pixel := img.At(x, y)
				// convert the pixel to RGBA
				_r, _g, _b, _a := pixel.RGBA()
				// get the values of the pixel
				r += _r
				g += _g
				b += _b
				a += _a
			}
			// get the average value of each pixel

			r = r / uint32(len(images))
			g = g / uint32(len(images))
			b = b / uint32(len(images))
			a = a / uint32(len(images))
			// set the pixel at x, y

			newImage.Set(x, y, &color.RGBA{
				R: uint8(r >> 8),
				G: uint8(g >> 8),
				B: uint8(b >> 8),
				A: uint8(a >> 8),
			})
		}
	}

	// save the new image
	err := saveImage(newImage, output)
	if err != nil {
		return fmt.Errorf("error saving image | %w", err)
	}

	return nil
}

func saveImage(img image.Image, output string) error {
	// open the output file
	outputFile, err := os.Create(output)
	if err != nil {
		return fmt.Errorf("error creating output file | %w", err)
	}
	defer outputFile.Close()

	// encode the image to the output file
	fileMimeType := mime.TypeByExtension(path.Ext(output))
	if fileMimeType == "image/jpeg" {
		err = jpeg.Encode(outputFile, img, nil)
		if err != nil {
			return fmt.Errorf("error encoding resized jpeg image | %w", err)
		}
	}

	if fileMimeType == "image/png" {
		err = png.Encode(outputFile, img)
		if err != nil {
			return fmt.Errorf("error encoding resized png image | %w", err)
		}
	}

	if fileMimeType == "image/tiff" {
		err = tiff.Encode(outputFile, img, nil)
		if err != nil {
			return fmt.Errorf("error encoding resized tiff image | %w", err)
		}
	}

	return nil
}

func getListOfFiles(args []string) ([]string, error) {
	var filenames []string
	patterns := args

	if len(patterns) == 0 {
		fmt.Println("No files provided, using stdin.")
		// read from stdin
		var err error
		fileLocationBytes, err := io.ReadAll(os.Stdin)
		if err != nil {
			return nil, fmt.Errorf("error reading from stdin | %w", err)
		}
		patterns = strings.Split(string(fileLocationBytes), "\n")
	}

	for _, pattern := range patterns {
		subFileNames, err := filepath.Glob(pattern)
		if err != nil {
			return nil, fmt.Errorf("error getting list of files | %w", err)
		}
		filenames = append(filenames, subFileNames...)
	}

	return filenames, nil
}