// Copyright (C) 2018 Marius Schellenberger

package main

import (
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"image"
	"image/color"
	_ "image/gif"
	_ "image/jpeg"
	"image/png"
	"io/ioutil"
	"os"
)

const offset = 4

func encodeLen(l int) []byte {
	b := make([]byte, offset)
	binary.LittleEndian.PutUint32(b, uint32(l))
	return b
}

func getLen(p image.Image) int {
	z := make([]byte, offset)
	for i := 0; i < offset; i++ {
		pix := color.NRGBAModel.Convert(p.At(i, 0)).(color.NRGBA)
		z[i] = (pix.R << 6)
		z[i] |= (pix.G << 6) >> 2
		z[i] |= (pix.B << 6) >> 4
		z[i] |= (pix.A << 6) >> 6
	}
	return int(binary.LittleEndian.Uint32(z))
}

func unseal(p image.Image) (b []byte) {
	rect := p.Bounds()
	l := getLen(p)
	c := 0
	for i := 0; i < rect.Max.Y; i++ {
		for j := 0; j < rect.Max.X; j++ {
			if i == 0 && j == 0 {
				j = offset
			}
			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 c == l {
				return
			}
			b = append(b, z)
			c++
		}
	}
	return
}

func seal(p image.Image, s []byte) (image.Image, error) {
	rect := p.Bounds()
	n := image.NewNRGBA(rect)
	c := 0
	arr := append(encodeLen(len(s)), s...)
	l := len(arr)
	max := (rect.Max.Y * rect.Max.X)
	if l > 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 < l {
				r := (pix.R &^ 3) | (arr[c] >> 6)
				g := (pix.G &^ 3) | (arr[c]<<2)>>6
				b := (pix.B &^ 3) | (arr[c]<<4)>>6
				a := (pix.A &^ 3) | (arr[c]<<6)>>6
				n.Set(j, i, color.NRGBA{r, g, b, a})
				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 file 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)
		}
	}
	if outFile != "-" {
		out, err = os.OpenFile(outFile, os.O_RDWR|os.O_CREATE, 0644)
		if err != nil {
			errf(err)
		}
	}
	p, _, err := image.Decode(in)
	if err != nil {
		errf(err)
	}
	if secret != "" {
		b, err := ioutil.ReadFile(secret)
		if err != nil {
			errf(err)
		}
		n, err := seal(p, b)
		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)
	}
}