package main

import (
	"flag"
	"fmt"
	"image"
	"image/color"
	_ "image/jpeg"
	"image/png"
	_ "image/png"
	"os"

	"slices"
)

func paletteFromImage(img image.Image) (p []color.Color) {
	for i := 0; i < img.Bounds().Dx(); i++ {
		p = append(p, img.At(i, 0))
	}
	return
}

func parseHexColor(s string) (c color.RGBA, err error) {
	c.A = 0xff
	switch len(s) {
	case 7:
		_, err = fmt.Sscanf(s, "#%02x%02x%02x", &c.R, &c.G, &c.B)
	case 4:
		_, err = fmt.Sscanf(s, "#%1x%1x%1x", &c.R, &c.G, &c.B)
		// Double the hex digits:
		c.R *= 17
		c.G *= 17
		c.B *= 17
	default:
		err = fmt.Errorf("invalid length, must be 7 or 4")

	}
	return
}

func mean(c1, c2 color.Color) color.Color {
	r1, g1, b1, _ := c1.RGBA()
	r2, g2, b2, _ := c2.RGBA()
	return color.RGBA{
		uint8((r1 + r2) / 2),
		uint8((g1 + g2) / 2),
		uint8((b1 + b2) / 2),
		0xFF,
	}
}

type Convert struct {
	image.Image
	color.Palette
}

func (c Convert) ColorModel() color.Model { return c.Image.ColorModel() }
func (c Convert) Bounds() image.Rectangle { return c.Image.Bounds() }
func (c Convert) At(x, y int) color.Color {
	if isMeanCloser(c.Palette, c.Image.At(x, y)) {
		first, second := twoClosest(c.Palette, c.Image.At(x, y))
		// chessboard texture
		if (x%2 != 0) && (y%2 != 0) || (x%2 == 0) && (y%2 == 0) {
			return first
		}
		return second
	}
	return c.Palette.Convert(c.Image.At(x, y))
}

func twoClosest(p color.Palette, c color.Color) (first, second color.Color) {
	first = p.Convert(c)
	i := p.Index(c)
	p2 := slices.Delete(slices.Clone(p), i, i+1)
	second = p2.Convert(c)
	return
}

func isMeanCloser(p color.Palette, c color.Color) bool {
	first, second := twoClosest(p, c)
	mp := color.Palette{first, mean(first, second), second}
	i := mp.Index(c)
	return i == mp.Index(mean(first, second))
}

func main() {
	paletteFile := flag.String("p", "", "provide palette file")
	flag.Parse()

	var palette []color.Color

	var colors []string
	_ = []string{
		"#ffffec", // soft-white
		"#ff0000", // red
		"#00ff00", // green
		"#0000ff", // blue
		"#000000", // black
	}
	funniPalette := []string{
		"#ffffec",
		"#000000",
		"#4a3544",
		"#757575",
		"#8e8e8e",
		"#666666",
		"#4c4c4c",
		"#3b3b3b",
		"#949494",
		"#ffffff",
		"#665361",
		"#3f3f3f",
		"#525252",
		"#838383",
		"#111111",
		"#070707",
		"#b9b9b9",
		"#dfdfdf",
	}
	colors = funniPalette
	if *paletteFile != "" {
		f, err := os.Open(*paletteFile)
		if err != nil {
			panic(err)
		}
		img, _, err := image.Decode(f)
		if err != nil {
			panic(err)
		}
		palette = paletteFromImage(img)
	} else {
		palette = func() (p []color.Color) {
			for _, v := range colors {
				c, err := parseHexColor(v)
				if err != nil {
					panic(err)
				}
				p = append(p, c)
			}
			return
		}()
	}
	img, _, err := image.Decode(os.Stdin)
	if err != nil {
		panic(err)
	}
	palette = append(palette, color.RGBA{88, 88, 88, 0xFF})
	err = png.Encode(os.Stdout, Convert{img, palette})
	if err != nil {
		panic(err)
	}
}