// Copyright (C) 2018 Marius Schellenberger

package main

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

func unseal(p image.Image) (b []byte) {
	rect := p.Bounds()
	for i := 0; i < rect.Max.Y; i++ {
		for j := 0; j < rect.Max.X; j++ {
			pix := color.NRGBAModel.Convert(p.At(j, i)).(color.NRGBA)
			z := (pix.R << 6)
			z |= (pix.G << 6) >> 2
			z |= (pix.B << 6) >> 4
			z |= (pix.A << 6) >> 6
			if z == byte(0) {
				b = append(b, '\n')
				return
			}
			b = append(b, z)
		}
	}
	return
}

func seal(p image.Image, s []byte) (image.Image, error) {
	rect := p.Bounds()
	n := image.NewNRGBA(rect)
	c := 0
	max := (rect.Max.Y * rect.Max.X) - 1
	if len(s) > max {
		return nil, errors.New("not enough space in image")
	}
	for i := 0; i < rect.Max.Y; i++ {
		for j := 0; j < rect.Max.X; j++ {
			pix := color.NRGBAModel.Convert(p.At(j, i)).(color.NRGBA)
			if c < len(s) {
				r := (pix.R &^ 3) | (s[c] >> 6)
				g := (pix.G &^ 3) | (s[c]<<2)>>6
				b := (pix.B &^ 3) | (s[c]<<4)>>6
				a := (pix.A &^ 3) | (s[c]<<6)>>6
				n.Set(j, i, color.NRGBA{r, g, b, a})
				c++
			} else if c == len(s) {
				n.Set(j, i, color.NRGBA{
					pix.R &^ 3,
					pix.G &^ 3,
					pix.B &^ 3,
					pix.A &^ 3,
				})
				c++
			} else {
				n.Set(j, i, pix)
			}
		}
	}
	return n, nil
}

var (
	inFile  string
	outFile string
	secret  string
)

func errf(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

func main() {
	flag.StringVar(&inFile, "in", "-", "input file (default: stdin)")
	flag.StringVar(&outFile, "out", "-", "output file, always png (default: stdout)")
	flag.StringVar(&secret, "secret", "", "secret to hide")
	flag.Parse()
	var (
		err error
		in  = os.Stdin
		out = os.Stdout
	)
	if inFile != "-" {
		in, err = os.Open(inFile)
		if err != nil {
			errf(err)
		}
	}
	p, _, err := image.Decode(in)
	if err != nil {
		errf(err)
	}
	if secret != "" {
		n, err := seal(p, []byte(secret))
		if err != nil {
			errf(err)
		}
		err = png.Encode(out, n)
		if err != nil {
			errf(err)
		}
		return
	}
	_, err = out.Write(unseal(p))
	if err != nil {
		errf(err)
	}
}
