package lafsession

import (
	"crypto/cipher"
	"encoding/binary"

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

const (
	recordHeaderLen = 1 + 8

	recordTypeData  byte = 0x10
	recordTypeClose byte = 0x11
)

func newAEAD(key []byte) (cipher.AEAD, error) {
	return chacha20poly1305.New(key)
}

func sealRecord(aead cipher.AEAD, seq uint64, recType byte, plaintext, dst []byte, nonce *[nonceLen]byte, aad *[recordHeaderLen]byte) []byte {
	aad[0] = recType
	binary.BigEndian.PutUint64(aad[1:], seq)

	for i := 0; i < 4; i++ {
		nonce[i] = 0
	}
	binary.BigEndian.PutUint64(nonce[4:], seq)

	need := recordHeaderLen + len(plaintext) + aead.Overhead()
	if cap(dst) < need {
		dst = make([]byte, need)
	} else {
		dst = dst[:need]
	}
	dst[0] = recType
	binary.BigEndian.PutUint64(dst[1:9], seq)
	_ = aead.Seal(dst[recordHeaderLen:recordHeaderLen], nonce[:], plaintext, aad[:])
	return dst
}

// openRecord decrypts a full frame. Returns record type and plaintext (reusing dst when possible).
func openRecord(aead cipher.AEAD, frame, dst []byte, nonce *[nonceLen]byte, aad *[recordHeaderLen]byte) (recType byte, plaintext []byte, err error) {
	if len(frame) < recordHeaderLen+aead.Overhead() {
		return 0, nil, ErrShortMessage
	}
	recType = frame[0]
	seq := binary.BigEndian.Uint64(frame[1:9])

	aad[0] = recType
	binary.BigEndian.PutUint64(aad[1:], seq)

	for i := 0; i < 4; i++ {
		nonce[i] = 0
	}
	binary.BigEndian.PutUint64(nonce[4:], seq)

	ct := frame[recordHeaderLen:]
	ptLen := len(ct) - aead.Overhead()
	if ptLen < 0 {
		return 0, nil, ErrShortMessage
	}
	if cap(dst) < ptLen {
		dst = make([]byte, ptLen)
	} else {
		dst = dst[:ptLen]
	}
	out, err := aead.Open(dst[:0], nonce[:], ct, aad[:])
	if err != nil {
		return 0, nil, ErrBadTag
	}
	return recType, out, nil
}