package main

import (
	"flag"
	"fmt"
	"image"
	"image/color"
	_ "image/jpeg"
	"image/png"
	"math/rand"
	"os"
	"time"
)

type coord struct {
	x, y int
}

// provided x, y, and color at location, return a color
type quantizerFunction func(int, int, color.Color) color.Color

// apply sequentially applies a quantizing function to an image and returns the result
func apply(i image.Image, f quantizerFunction) image.Image {
	out := image.NewRGBA(image.Rect(0, 0, i.Bounds().Max.X, i.Bounds().Max.Y))

	b := out.Bounds()
	for y := b.Min.Y; y < b.Max.Y; y++ {
		for x := b.Min.X; x < b.Max.X; x++ {
			out.Set(x, y, f(x, y, i.At(x, y)))
		}
	}

	return out
}

type ditherOpts struct {
	raw    string
	level  int
	invert bool
}

func (o *ditherOpts) Set(input string) error {
	for _, p := range input {
		switch p {
		case 'n':
			o.invert = true
		case '0':
			o.level = 0
		case '1':
			o.level = 1
		default:
			return fmt.Errorf("unknown option %s", string(p))
		}
	}
	return nil
}

func (o *ditherOpts) String() string {
	return o.raw
}

func main() {
	ditherOpts := &ditherOpts{}

	filenameFlag := flag.String("in", "", "Input filename. PNG or JPG formats accepted.")
	verboseFlag := flag.Bool("v", false, "Include diagnostic log lines.")
	paletteFlag := flag.String("p", "bw",
		`Palette to use. Defaults to 'bw': black-and-white.
Pass 'sixteen' for the classic sixteen-color palette.
Not all dither options support palettes; unsupported ditherers will ignore this flag.`)
	dithererFlag := flag.String("d", "naive",
		`Ditherer to use. Defaults to 'naive', which just flattens to nearest color.
Other options: 'noise', 'bayer', 'error', 'floydsteinberg', 'jjn', 'atkinson'.`)
	flag.Var(ditherOpts, "o",
		`Supply options for the selected ditherer.
Currently only relevant to the bayer ditherer, which defaults to 'level 0, non-inverted'.
To change: '1' applies a level-one bayer diffuser; '1n' would invert it, '0' would be the default.`)
	outputFlag := flag.String("out", "output.png", "Output filename. Outputs in PNG format.")

	flag.Parse()

	// extract all those pointers
	filename := *filenameFlag
	ditherer := *dithererFlag
	palette := *paletteFlag
	outfile := *outputFlag
	verbose := *verboseFlag

	if filename == "" || outfile == "" {
		fmt.Printf("Usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
		os.Exit(0)
	}

	i, codex, err := loadImage(filename)
	if err != nil {
		fmt.Printf("error loading %s; %v\n", filename, err)
		os.Exit(255)
	}
	if verbose {
		fmt.Printf("loaded %s using %s\n", filename, codex)
	}

	t := time.Now()

	var new image.Image
	switch ditherer {
	case "noop":
		new = apply(i, noOp)
	case "naive":
		if palette == "bw" {
			new = apply(i, naiveBW)
		} else if palette == "sixteen" {
			new = apply(i, naivePalette(sixteencolors))
		}
	case "noise":
		rand.Seed(time.Now().UnixNano())
		if palette == "bw" {
			new = apply(i, randomNoise)
		} else if palette == "sixteen" {
			new = apply(i, randomNoisePalette(sixteencolors))
		} else {
			fmt.Printf("unknown palette option for %s ditherer: %s\n", ditherer, palette)
			os.Exit(2)
		}
	case "bayer":
		if palette == "bw" {
			new = apply(i, bayerDithering(ditherOpts.level, ditherOpts.invert))
		} else if palette == "sixteen" {
			new = apply(i, colorBayer(ditherOpts.level, sixteencolors))
		} else {
			fmt.Printf("unknown palette option for %s ditherer: %s\n", ditherer, palette)
			os.Exit(2)
		}
	case "error":
		if palette == "bw" {
			new = apply(i, simpleErrorDiffusion())
		} else if palette == "sixteen" {
			new = apply(i, simpleColorErrorDiffusion())
		} else {
			fmt.Printf("unknown palette option for %s ditherer: %s\n", ditherer, palette)
			os.Exit(2)
		}
	case "floydsteinberg":
		new = apply(i, floydSteinberg())
	case "jjn":
		if palette == "bw" {
			new = apply(i, jarvisJudiceNinke())
		} else if palette == "sixteen" {
			new = apply(i, colorJarvisJudiceNinke())
		} else {
			fmt.Printf("unknown palette option for %s ditherer: %s\n", ditherer, palette)
			os.Exit(2)
		}
	case "atkinson":
		new = apply(i, atkinson())
	default:
		fmt.Printf("unknown ditherer option: %s\n", ditherer)
		os.Exit(2)
	}

	if verbose {
		fmt.Printf("dithering took %s\n", time.Since(t))
	}

	err = saveImage(outfile, new)
	if err != nil {
		fmt.Printf("error saving %s; %v\n", outfile, err)
		os.Exit(255)
	}

	if verbose {
		fmt.Printf("saved output to %s\n", outfile)
	}
}

func loadImage(filename string) (image.Image, string, error) {
	r, err := os.Open(filename)
	if err != nil {
		return nil, "", err
	}
	return image.Decode(r)
}

func saveImage(filename string, img image.Image) error {
	f, err := os.Create(filename)
	if err != nil {
		return err
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		return err
	}

	if err := f.Close(); err != nil {
		return err
	}

	return nil
}