// Package chacha20 implements the ChaCha20 stream cipher as specified in
// RFC 8439.
//
// The block counter starts at zero, as required by NIP-44. The package is
// pure Go and uses no architecture-specific code, so it builds with TinyGo.
package chacha20

import (
	"encoding/binary"
	"math/bits"
)

const (
	// KeySize is the length of a ChaCha20 key in bytes.
	KeySize = 32
	// NonceSize is the length of an RFC 8439 nonce in bytes.
	NonceSize = 12
	// BlockSize is the number of keystream bytes produced per block.
	BlockSize = 64
)

// sigma holds the ChaCha20 constant "expand 32-byte k", one word per
// little-endian byte quadruple.
var sigma = [4]uint32{0x61707865, 0x3320646e, 0x79622d32, 0x6b206574}

// quarter performs one ChaCha20 quarter round on four state words.
func quarter(a, b, c, d uint32) (uint32, uint32, uint32, uint32) {
	a += b
	d ^= a
	d = bits.RotateLeft32(d, 16)
	c += d
	b ^= c
	b = bits.RotateLeft32(b, 12)
	a += b
	d ^= a
	d = bits.RotateLeft32(d, 8)
	c += d
	b ^= c
	b = bits.RotateLeft32(b, 7)
	return a, b, c, d
}

// initState builds the initial ChaCha20 state for the given key, counter and
// nonce.
func initState(key *[KeySize]byte, counter uint32, nonce *[NonceSize]byte) [16]uint32 {
	var s [16]uint32
	s[0], s[1], s[2], s[3] = sigma[0], sigma[1], sigma[2], sigma[3]
	for i := 0; i < 8; i++ {
		s[4+i] = binary.LittleEndian.Uint32(key[4*i : 4*i+4])
	}
	s[12] = counter
	for i := 0; i < 3; i++ {
		s[13+i] = binary.LittleEndian.Uint32(nonce[4*i : 4*i+4])
	}
	return s
}

// Block returns the 64-byte ChaCha20 keystream block for the given key,
// block counter and nonce.
func Block(key *[KeySize]byte, counter uint32, nonce *[NonceSize]byte) [BlockSize]byte {
	state := initState(key, counter, nonce)
	x := state

	// Ten double rounds: column rounds followed by diagonal rounds.
	for i := 0; i < 10; i++ {
		x[0], x[4], x[8], x[12] = quarter(x[0], x[4], x[8], x[12])
		x[1], x[5], x[9], x[13] = quarter(x[1], x[5], x[9], x[13])
		x[2], x[6], x[10], x[14] = quarter(x[2], x[6], x[10], x[14])
		x[3], x[7], x[11], x[15] = quarter(x[3], x[7], x[11], x[15])
		x[0], x[5], x[10], x[15] = quarter(x[0], x[5], x[10], x[15])
		x[1], x[6], x[11], x[12] = quarter(x[1], x[6], x[11], x[12])
		x[2], x[7], x[8], x[13] = quarter(x[2], x[7], x[8], x[13])
		x[3], x[4], x[9], x[14] = quarter(x[3], x[4], x[9], x[14])
	}

	var out [BlockSize]byte
	for i := 0; i < 16; i++ {
		binary.LittleEndian.PutUint32(out[4*i:4*i+4], x[i]+state[i])
	}
	return out
}

// XORKeyStream XORs each byte of src with the ChaCha20 keystream derived from
// key and nonce, writing the result to dst. The block counter starts at zero
// and is incremented for every 64-byte block.
//
// dst must be at least as long as src. dst and src may be the same slice, or
// may be non-overlapping; they must not partially overlap.
func XORKeyStream(key *[KeySize]byte, nonce *[NonceSize]byte, dst, src []byte) {
	if len(dst) < len(src) {
		panic("chacha20: output smaller than input")
	}
	var counter uint32
	for len(src) > 0 {
		ks := Block(key, counter, nonce)
		counter++
		n := len(src)
		if n > BlockSize {
			n = BlockSize
		}
		for i := 0; i < n; i++ {
			dst[i] = src[i] ^ ks[i]
		}
		dst = dst[n:]
		src = src[n:]
	}
}