package symm

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

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

var ErrInvalidChaCha20NonceLength = errors.New("chacha20 nonce length must be 12 or 24 bytes")

func EncryptChaCha20(data, key, nonce []byte) ([]byte, error) {
	stream, err := chacha20.NewUnauthenticatedCipher(key, nonce)
	if err != nil {
		return nil, err
	}
	out := make([]byte, len(data))
	stream.XORKeyStream(out, data)
	return out, nil
}

func DecryptChaCha20(src, key, nonce []byte) ([]byte, error) {
	return EncryptChaCha20(src, key, nonce)
}

func EncryptChaCha20Stream(dst io.Writer, src io.Reader, key, nonce []byte) error {
	stream, err := chacha20.NewUnauthenticatedCipher(key, nonce)
	if err != nil {
		return err
	}
	return xorStreamCopy(dst, src, stream)
}

func DecryptChaCha20Stream(dst io.Writer, src io.Reader, key, nonce []byte) error {
	return EncryptChaCha20Stream(dst, src, key, nonce)
}

func EncryptChaCha20Poly1305(plain, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newChaCha20AEAD(key, nonce)
	if err != nil {
		return nil, err
	}
	return aead.Seal(nil, nonce, plain, aad), nil
}

func DecryptChaCha20Poly1305(ciphertext, key, nonce, aad []byte) ([]byte, error) {
	aead, err := newChaCha20AEAD(key, nonce)
	if err != nil {
		return nil, err
	}
	return aead.Open(nil, nonce, ciphertext, aad)
}

func newChaCha20AEAD(key, nonce []byte) (cipher.AEAD, error) {
	switch len(nonce) {
	case chacha20poly1305.NonceSize:
		return chacha20poly1305.New(key)
	case chacha20poly1305.NonceSizeX:
		return chacha20poly1305.NewX(key)
	default:
		return nil, ErrInvalidChaCha20NonceLength
	}
}