package crypto

import (
	"crypto/cipher"
	"errors"
	"io"

	"golang.org/x/crypto/chacha20poly1305"
)

type cryptoReader interface {
	io.Reader
}

type nonce struct {
	N [NonceSize]byte
	C int
}

type reader struct {
	a   cipher.AEAD
	src cryptoReader

	unread []byte // decrypted but unread data, backed by buf
	buf    [FullSize]byte
	off    int64

	err    error
	nonce  [NonceSize]byte
	cn     int64
	outBuf []byte
}

func newReader(key []byte, src cryptoReader) (*reader, error) {
	aead, err := chacha20poly1305.New(key)
	if err != nil {
		return nil, err
	}
	return &reader{
		a:      aead,
		src:    src,
		cn:     -1,
		outBuf: make([]byte, 0, FullSize),
	}, nil
}

func (r *reader) Read(p []byte) (int, error) {
	if len(r.unread) > 0 {
		n := copy(p, r.unread)
		r.unread = r.unread[n:]
		return n, nil
	}
	if r.err != nil {
		return 0, r.err
	}
	if len(p) == 0 {
		return 0, nil
	}

	last, err := r.readChunk()
	if err != nil {
		r.err = err
		return 0, err
	}
	if r.off > 0 && len(r.unread) >= int(r.off) {
		r.unread = r.unread[r.off:]
		r.off = 0
	}

	n := copy(p, r.unread)
	r.unread = r.unread[n:]

	if last {
		r.err = io.EOF
	}

	return n, nil
}

func (r *reader) readChunk() (last bool, err error) {
	if len(r.unread) != 0 {
		panic("stream: internal error: readChunk called with dirty buffer")
	}

	in := r.buf[:]
	n, err := io.ReadFull(r.src, in)
	switch {
	case err == io.EOF:
		// A message can't end without a marked chunk. This message is truncated.
		return false, io.ErrUnexpectedEOF
	case err == io.ErrUnexpectedEOF:
		// The last chunk can be short.
		in = in[:n]
		last = true
		setLastChunkFlag(&r.nonce)
	case err != nil:
		return false, err
	}

	out, err := r.a.Open(r.outBuf, r.nonce[:], in, nil)
	if err != nil && !last {
		// Check if this was a full-length final chunk.
		last = true
		setLastChunkFlag(&r.nonce)
		out, err = r.a.Open(r.outBuf, r.nonce[:], in, nil)
	}
	if err != nil {
		return false, errors.New("failed to decrypt and authenticate payload chunk")
	}

	incNonce(&r.nonce)
	size := FullSize
	if len(out) < FullSize {
		size = len(out)
	}
	r.unread = out[:size]
	return last, nil
}