package cryptoutil

import (
	"bufio"
	"bytes"
	"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
)

// 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           int64
	noncePrefix [noncePrefixSize]byte
}

func (o *oae2) initialize(header []byte) error {
	o.i = 1
	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:], uint64(o.i))
	if lastBlock {
		nonce[gcmNonceSize-1] = 1
	}
	return nil
}

func (o *oae2) encryptBlock(out, 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(out, nonce, block, o.additionalData)
	o.i++
	o.additionalData = nil
	return encrypted, nil
}

func (o *oae2) decryptBlock(out, 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(out, nonce, block, o.additionalData)
	if err != nil {
		return nil, err
	}
	o.i++
	o.additionalData = nil
	return decrypted, nil
}

// 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
	segmentSize int

	initialized bool
	err         error
	buf         []byte
}

type opts struct {
	additionalData []byte
	segmentSize    int
}

var defaultOpts = opts{segmentSize: 192*1024 - 1 - aeadOverhead}

type Option func(*opts)

func WithAdditionalData(ad []byte) Option {
	return func(opts *opts) {
		opts.additionalData = ad
	}
}

func WithSegmentSize(segmentSize int) Option {
	return func(opts *opts) {
		opts.segmentSize = segmentSize
	}
}

// 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, options ...Option) *EncryptingWriter {
	opts := defaultOpts
	for _, o := range options {
		o(&opts)
	}
	return &EncryptingWriter{
		w: w,
		oae2: oae2{
			key:            k,
			additionalData: opts.additionalData,
		},
		segmentSize: opts.segmentSize,
	}
}

func (w *EncryptingWriter) initialize() error {
	w.initialized = true
	w.buf = make([]byte, 0, w.segmentSize+aeadOverhead)
	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
	}
	r := bytes.NewReader(buf)
	nn := 0
	for r.Len() > 0 {
		if len(w.buf) == w.segmentSize {
			var encrypted []byte
			if encrypted, w.err = w.oae2.encryptBlock(w.buf[:0], w.buf, false); w.err != nil {
				return nn, w.err
			}
			if _, w.err = w.w.Write(encrypted); w.err != nil {
				return nn, w.err
			}
			w.buf = w.buf[:0]
		}
		n, _ := r.Read(w.buf[len(w.buf):w.segmentSize])
		w.buf = w.buf[:len(w.buf)+n]
		nn += 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[:0], w.buf, 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           *bufio.Reader
	oae2        oae2
	segmentSize int

	initialized  bool
	decryptedBuf bytes.Buffer
}

// 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, options ...Option) *DecryptingReader {
	opts := defaultOpts
	for _, o := range options {
		o(&opts)
	}
	return &DecryptingReader{
		r: bufio.NewReaderSize(r, opts.segmentSize+aeadOverhead+1),
		oae2: oae2{
			key:            k,
			additionalData: opts.additionalData,
		},
		segmentSize: opts.segmentSize,
	}
}

func (r *DecryptingReader) initialize() error {
	header, err := r.r.Peek(headerSize)
	if err != nil {
		return err
	}
	if err := r.oae2.initialize(header); err != nil {
		return err
	}
	r.r.Discard(len(header))
	r.decryptedBuf = *bytes.NewBuffer(make([]byte, 0, r.segmentSize))
	r.initialized = true
	return nil
}

func (r *DecryptingReader) fillBuf() error {
	encryptedSegmentSize := r.segmentSize + aeadOverhead
	// Peek one extra byte to make sure if this is the last segment
	block, readErr := r.r.Peek(encryptedSegmentSize + 1)
	if len(block) == 0 {
		return readErr
	}
	block = block[:min(len(block), encryptedSegmentSize)]
	r.decryptedBuf.Reset()
	result, err := r.oae2.decryptBlock(r.decryptedBuf.AvailableBuffer(), block, readErr == io.EOF)
	if err != nil {
		return err
	}
	r.decryptedBuf.Write(result)
	r.r.Discard(len(block))
	return nil
}

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