package rabaead

import (
	"bytes"
	"crypto/cipher"
	"encoding/binary"
	"errors"
	"io"

	"snix.ir/rabbitio"
)

const cmrs = 0x02 // chunk size indicator,
// without this reader cannot calculate actual size of plaintext

// additional data func, return value is used as AD in Seal and Open
// nil AdFunc is harmless and equal to func()[]byte{return nil}
type AdditionalFunc func() []byte

type chunkReader struct {
	aead  cipher.AEAD
	csize int
	rader io.Reader
	buff  []byte
	nonce []byte
	adexe AdditionalFunc
}

type chunkWriter struct {
	aead   cipher.AEAD
	csize  int
	writer io.Writer
	buff   []byte
	nonce  []byte
	adexe  AdditionalFunc
}

// NewChunkReader returns a chunkReader data type, this reader reads and open() aead
// ciphertext, each chunk has its own tag and cmrk value.
// this reader has a chunk size in-memory buffer, large chunk size can make application to runs
// out of memory, thus is most suitable for sliced data, like network data transmit and so..
func NewChunkReader(r io.Reader, chnk int, a cipher.AEAD, nonce []byte, f AdditionalFunc) (*chunkReader, error) {

	if len(nonce) != rabbitio.IVXLen && len(nonce) != 0 {
		return nil, rabbitio.ErrInvalidIVX
	}

	if chnk > int(^uint16(0)) || chnk <= 0 {
		return nil, errors.New("rabaead: bad chunk size")
	}

	s := &chunkReader{
		aead:  a,
		buff:  []byte{},
		nonce: make([]byte, len(nonce)),
		csize: chnk,
		rader: r,
		adexe: f,
	}

	if s.adexe == nil {
		s.adexe = func() []byte { return nil }
	}
	copy(s.nonce, nonce)
	return s, nil
}

// NewChunkWriter returns a chunkWriter data type, this writer sale() and write aead
// plaintext, each chunk has its own tag and cmrk value.
// this writer has a chunk size in-memory buffer, large chunk size can make application to
// runs out of memory, thus is most suitable for sliced data, like network data transmit and so..
func NewChunkWriter(w io.Writer, chnk int, a cipher.AEAD, nonce []byte, f AdditionalFunc) (*chunkWriter, error) {

	if len(nonce) != rabbitio.IVXLen && len(nonce) != 0 {
		return nil, rabbitio.ErrInvalidIVX
	}

	if chnk > int(^uint16(0)) || chnk <= 0 {
		return nil, errors.New("rabaead: bad chunk size")
	}

	s := &chunkWriter{
		aead:   a,
		buff:   []byte{},
		nonce:  make([]byte, len(nonce)),
		csize:  chnk,
		writer: w,
		adexe:  f,
	}

	if s.adexe == nil {
		s.adexe = func() []byte { return nil }
	}

	copy(s.nonce, nonce)
	return s, nil
}

// Close method, if there is any
func (w *chunkWriter) Close() error {
	if c, ok := w.writer.(io.Closer); ok {
		return c.Close()
	}
	return nil
}

// Write writes plaintext chunk into the sale() and underlying writer
// write would not report overhead data (chunk size marker and poly1305 tag) in
// written return value. for each chunk there is 2+16 byte overhead data.
// AdFunc will be triggered for each chunk of data
func (w *chunkWriter) Write(b []byte) (n int, err error) {
	w.buff = b
	for len(w.buff) > 0 {
		s, err := w.write()
		if err != nil {
			return n, err
		}
		n += s
	}
	return
}

func (w *chunkWriter) write() (int, error) {
	size := cmrs + w.csize + w.aead.Overhead()
	chnk := make([]byte, size)
	var n int
	var err error

	if len(w.buff) > 0 {
		s := copy(chnk[cmrs:len(chnk)-w.aead.Overhead()], w.buff)
		w.buff = w.buff[s:]
		copy(chnk[0:cmrs], uint16Little(uint16(s)))

		w.aead.Seal(chnk[:0], w.nonce, chnk[:cmrs+w.csize], w.adexe())
		_, err = w.writer.Write(chnk)
		if err != nil {
			return n, err
		}
		n += s
	}

	return n, err
}

// Read reads and open() ciphertext chunk from underlying reader
// read would not report overhead data (chunk size marker and poly1305 tag) in its
// return value. if the read data from underlying reader is corrupted, ErrAuthMsg
// error will be returned. for each chunk there is 2+16 byte overhead data.
// AdFunc will be triggered for each chunk of data

func (r *chunkReader) Read(b []byte) (int, error) {
	if len(b) <= r.csize {
		return r.readTo(b)
	}
	n := 0
	for {
		if n+r.csize > len(b) {
			sr, err := r.readTo(b[n:])
			n += sr
			if err != nil {
				return n, err
			}
			break
		}
		sr, err := r.readTo(b[n : n+r.csize])
		n += sr
		if err != nil {
			return n, err
		}
	}

	return n, nil
}

func (r *chunkReader) readTo(b []byte) (int, error) {
	var n int
	if len(r.buff) > 0 {
		n = copy(b, r.buff)
		r.buff = r.buff[n:]
		return n, nil
	}

	sr, err := r.read()
	n = copy(b, r.buff[:sr])
	r.buff = r.buff[n:]
	return n, err
}

func (r *chunkReader) read() (int, error) {

	var n int
	size := cmrs + r.csize + r.aead.Overhead()
	chnk := make([]byte, size)
	chLE := uint16Little(uint16(r.csize))

	si, err := io.ReadFull(r.rader, chnk)
	if err != nil {
		return n, err
	}

	if si > 0 {
		_, err = r.aead.Open(chnk[:0], r.nonce, chnk, r.adexe())
		if err != nil {
			return n, err
		}

		if bytes.Equal(chnk[0:cmrs], chLE) {
			n += r.csize
			r.buff = append(r.buff, chnk[cmrs:cmrs+r.csize]...)
		} else {
			f := binary.LittleEndian.Uint16(chnk[0:cmrs])
			n += int(f)
			r.buff = append(r.buff, chnk[cmrs:cmrs+f]...)
		}
	}

	return n, err
}

func uint16Little(n uint16) []byte {
	b := make([]byte, cmrs)
	binary.LittleEndian.PutUint16(b, n)
	return b
}