package cryptoutil

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hkdf"
	"crypto/rand"
	"crypto/sha512"
	"encoding/binary"
	"errors"
	"io"
)

// Online Authenticated Encryption from https://eprint.iacr.org/2015/189.pdf

const (
	aeadOverhead         = 16
	aesKeySize           = 32
	noncePrefixSize      = 3
	gcmNonceSize         = 12
	headerSize           = aesKeySize + noncePrefixSize
	cacheSize            = 192 * 1024
	encryptedSegmentSize = cacheSize - 1
	segmentSize          = encryptedSegmentSize - aeadOverhead
)

// An EncryptionKey is used for encrypting and decrypting data. A key can be any
// length, but at least 64 bytes would be recommended.
type EncryptionKey []byte

type oae2 struct {
	key            EncryptionKey
	additionalData []byte

	aead        cipher.AEAD
	i           uint64
	noncePrefix [noncePrefixSize]byte
}

func (o *oae2) initialize(header []byte) error {
	copy(o.noncePrefix[:], header[aesKeySize:])
	derivedKey, err := hkdf.Key(sha512.New, o.key, header[:aesKeySize], "", aesKeySize)
	if err != nil {
		return err
	}
	block, err := aes.NewCipher(derivedKey)
	if err != nil {
		return err
	}
	o.aead, err = cipher.NewGCM(block)
	return err
}

func (o *oae2) nonce(nonce []byte, lastBlock bool) error {
	copy(nonce, o.noncePrefix[:])
	if o.i == 0 {
		return errors.New("counter overflowed (64 bits??)")
	}
	binary.BigEndian.PutUint64(nonce[noncePrefixSize:], o.i)
	o.i++
	if lastBlock {
		nonce[gcmNonceSize-1] = 1
	}
	return nil
}

func (o *oae2) encryptBlock(block []byte, lastBlock bool) ([]byte, error) {
	nonce := make([]byte, gcmNonceSize)
	if err := o.nonce(nonce, lastBlock); err != nil {
		return nil, err
	}
	encrypted := o.aead.Seal(block[:len(block)+aeadOverhead][:0], nonce, block, o.additionalData)
	o.additionalData = nil
	return encrypted, nil
}

func (o *oae2) decryptBlock(block []byte, lastBlock bool) ([]byte, error) {
	nonce := make([]byte, gcmNonceSize)
	if err := o.nonce(nonce, lastBlock); err != nil {
		return nil, err
	}
	decrypted, err := o.aead.Open(block[:0], nonce, block, o.additionalData)
	o.additionalData = nil
	return decrypted, err
}

// An EncryptingWriter encrypts data in segments using the STREAM construction
// described in https://eprint.iacr.org/2015/189.pdf. The writer buffers data up
// to the segment size, so it's important to call Close to flush the
// final segment.
type EncryptingWriter struct {
	w    io.Writer
	oae2 oae2

	initialized bool
	err         error

	bufN int
	buf  [encryptedSegmentSize]byte
}

// NewWriter returns a new EncryptingWriter that writes to w. The additionalData
// will be authenticated with the first segment, but is not written to w. The
// same additional data must be provided when decrypting.
func (k EncryptionKey) NewWriter(w io.Writer, additionalData []byte) *EncryptingWriter {
	return &EncryptingWriter{
		w: w,
		oae2: oae2{
			key:            k,
			additionalData: additionalData,
			i:              1,
		},
	}
}

func (w *EncryptingWriter) initialize() error {
	w.initialized = true
	header := make([]byte, headerSize)
	rand.Read(header)
	if w.err = w.oae2.initialize(header); w.err != nil {
		return w.err
	}
	_, w.err = w.w.Write(header)
	return w.err
}

func (w *EncryptingWriter) Write(buf []byte) (int, error) {
	if !w.initialized {
		if err := w.initialize(); err != nil {
			return 0, err
		}
	}
	if w.err != nil {
		return 0, w.err
	}
	nn := 0
	for len(buf) > 0 {
		if w.bufN == segmentSize {
			var encrypted []byte
			if encrypted, w.err = w.oae2.encryptBlock(w.buf[:w.bufN], false); w.err != nil {
				return nn, w.err
			}
			if _, w.err = w.w.Write(encrypted); w.err != nil {
				return nn, w.err
			}
			w.bufN = 0
		}
		n := copy(w.buf[w.bufN:segmentSize], buf)
		w.bufN += n
		nn += n
		buf = buf[n:]
	}
	return nn, nil
}

var errClosed = errors.New("closed")

// Close encrypts writes the final segment. It does not close the
// underlying writer.
func (w *EncryptingWriter) Close() error {
	if !w.initialized {
		if err := w.initialize(); err != nil {
			return err
		}
	}
	if w.err == errClosed {
		return nil
	}
	if w.err != nil {
		return w.err
	}
	var encrypted []byte
	if encrypted, w.err = w.oae2.encryptBlock(w.buf[:w.bufN], true); w.err != nil {
		return w.err
	}
	if _, w.err = w.w.Write(encrypted); w.err != nil {
		return w.err
	}
	w.err = errClosed
	return nil
}

// A DecryptingReader decrypts data using the STREAM construction.
type DecryptingReader struct {
	r    io.Reader
	oae2 oae2

	initialized bool
	err         error

	bufRead, bufN int
	peekedByte    bool
	buf           [encryptedSegmentSize + 1]byte
}

// NewReader returns a new DecryptingWriter that decrypts data from r. The
// additionaData must be the same that was provided when encrypting.
func (k EncryptionKey) NewReader(r io.Reader, additionalData []byte) *DecryptingReader {
	return &DecryptingReader{
		r: r,
		oae2: oae2{
			key:            k,
			additionalData: additionalData,
			i:              1,
		},
	}
}

func (r *DecryptingReader) initialize() error {
	r.initialized = true
	header := make([]byte, headerSize)
	if _, r.err = io.ReadFull(r.r, header); r.err != nil {
		return r.err
	}
	r.err = r.oae2.initialize(header)
	return r.err
}

func (r *DecryptingReader) fillBuf() error {
	n := 0
	if r.peekedByte {
		n = 1
		r.buf[0] = r.buf[encryptedSegmentSize]
		r.peekedByte = false
	}
	for n < len(r.buf) && r.err == nil {
		var m int
		m, r.err = r.r.Read(r.buf[n:])
		n += m
	}
	if n == 0 {
		return r.err
	}
	if n == encryptedSegmentSize+1 {
		r.peekedByte = true
		n = encryptedSegmentSize
	}
	result, err := r.oae2.decryptBlock(r.buf[:n], r.err == io.EOF)
	if err != nil {
		r.err = err
		return err
	}
	r.bufRead = 0
	r.bufN = len(result)
	return nil
}

func (r *DecryptingReader) Read(buf []byte) (int, error) {
	if !r.initialized {
		if err := r.initialize(); err != nil {
			return 0, err
		}
	}
	if r.bufRead == r.bufN {
		if err := r.fillBuf(); err != nil {
			return 0, err
		}
	}
	n := copy(buf, r.buf[r.bufRead:r.bufN])
	r.bufRead += n
	if r.bufRead == r.bufN {
		return n, r.err
	}
	return n, nil
}