package rabaead

import (
	"crypto/cipher"
	"errors"

	"snix.ir/poly1305"
	"snix.ir/rabbitio"
	"snix.ir/rabbitio/subtle"
)

const polykeylen = 0x20 // poly1305 key len: 32byte

var ErrAuthMsg = errors.New("rabaead: message authentication failed")

type rabbitPoly1305 struct {
	key       []byte // rabbit cipher key
	noncesize int    // rabbit iv size
}

// NewAEAD returns a rabbit aead data-type
// key must be 16 byte len
func NewAEAD(key []byte) (cipher.AEAD, error) { return newRabbitAead(key) }

func newRabbitAead(key []byte) (cipher.AEAD, error) {
	if len(key) != rabbitio.KeyLen {
		return nil, rabbitio.ErrInvalidKey
	}

	rabbitAead := &rabbitPoly1305{
		noncesize: rabbitio.IVXLen,
		key:       make([]byte, rabbitio.KeyLen),
	}
	copy(rabbitAead.key, key)
	return rabbitAead, nil

}

// Overhead returns poly1305 tag size: 16byte
func (c *rabbitPoly1305) Overhead() int { return poly1305.TagSize }

// NonceSize returns rabbit iv len: 8byte
func (c *rabbitPoly1305) NonceSize() int { return c.noncesize }

func (c *rabbitPoly1305) sealRabbit(dst, nonce, plaintext, ad []byte) []byte {
	ret, out := headtail(dst, len(plaintext)+poly1305.TagSize)
	ciphertext, tag := out[:len(plaintext)], out[len(plaintext):]
	if subtle.InexactOverlap(out, plaintext) {
		panic("rabaead: invalid buffer memory overlap")
	}

	var polyKey [polykeylen]byte
	s, err := rabbitio.NewCipher(c.key, nonce)
	if err != nil {
		panic(err)
	}
	s.XORKeyStream(polyKey[:], polyKey[:])
	p := poly1305.New(&polyKey)
	writePadding(p, ad)

	s, err = rabbitio.NewCipher(c.key, nonce)
	if err != nil {
		panic(err)
	}
	s.XORKeyStream(ciphertext, plaintext)

	writePadding(p, ciphertext)

	writeUint64(p, len(ad))
	writeUint64(p, len(plaintext))
	p.Sum(tag[:0x00])

	return ret
}

func (c *rabbitPoly1305) openRabbit(dst, nonce, ciphertext, ad []byte) ([]byte, error) {
	tag := ciphertext[len(ciphertext)-poly1305.TagSize:]
	ciphertext = ciphertext[:len(ciphertext)-poly1305.TagSize]

	var polyKey [polykeylen]byte
	s, err := rabbitio.NewCipher(c.key, nonce)
	if err != nil {
		panic(err)
	}
	s.XORKeyStream(polyKey[:], polyKey[:])

	p := poly1305.New(&polyKey)
	writePadding(p, ad)
	writePadding(p, ciphertext)

	writeUint64(p, len(ad))
	writeUint64(p, len(ciphertext))

	ret, out := headtail(dst, len(ciphertext))
	if subtle.InexactOverlap(out, ciphertext) {
		panic("rabaead: invalid buffer memory overlap")
	}

	// check data integrity
	if !p.Verify(tag) {
		return nil, ErrAuthMsg
	}

	s, err = rabbitio.NewCipher(c.key, nonce)
	if err != nil {
		panic(err)
	}
	s.XORKeyStream(out, ciphertext)
	return ret, nil
}

// Open opens a rabbit aead ciphertext.
// panic occurs if nonce len is not equal to IVXLen (8byte) or zero
// if data is not verified, ErrAuthMsg will be returned
func (c *rabbitPoly1305) Open(dst, nonce, ciphertext, ad []byte) ([]byte, error) {
	if len(ciphertext) < poly1305.TagSize {
		return nil, ErrAuthMsg
	}

	if uint64(len(ciphertext)) > (1<<38)-48 {
		panic("rabaead: ciphertext too large")
	}

	return c.openRabbit(dst, nonce, ciphertext, ad)
}

// Seal seals a plaintext into the rabbit aead ciphertext.
// panic occurs if nonce len is not equal to IVXLen (8byte) or zero
func (c *rabbitPoly1305) Seal(dst, nonce, plaintext, ad []byte) []byte {
	if uint64(len(plaintext)) > (1<<38)-64 {
		panic("rabaead: plaintext too large")
	}

	return c.sealRabbit(dst, nonce, plaintext, ad)
}