package chacha8

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

// NetEase 4229938 UXChaCha8 stream cipher.
// Uses custom Sigma constants differing from standard ChaCha.
const (
	Sigma0 uint32 = 0x60707865
	Sigma1 uint32 = 0x3321646e
	Sigma2 uint32 = 0x79622d32
	Sigma3 uint32 = 0x6b206e74
)

type UxChaCha8 struct {
	state        [16]uint32
	keystream    [64]byte
	keystreamPos int
}

func New(key []byte, nonce []byte) (*UxChaCha8, error) {
	if len(key) != 32 {
		return nil, errors.New("UXChaCha8 key must be 32 bytes")
	}
	if len(nonce) != 12 {
		return nil, errors.New("UXChaCha8 nonce must be 12 bytes")
	}

	c := &UxChaCha8{
		keystreamPos: 64,
	}

	c.state[0] = Sigma0
	c.state[1] = Sigma1
	c.state[2] = Sigma2
	c.state[3] = Sigma3

	for i := 0; i < 8; i++ {
		c.state[4+i] = binary.LittleEndian.Uint32(key[i*4 : (i+1)*4])
	}

	c.state[12] = 0
	c.state[13] = binary.LittleEndian.Uint32(nonce[0:4])
	c.state[14] = binary.LittleEndian.Uint32(nonce[4:8])
	c.state[15] = binary.LittleEndian.Uint32(nonce[8:12])

	return c, nil
}

func (c *UxChaCha8) XorInPlace(data []byte) {
	offset := 0
	for offset < len(data) {
		if c.keystreamPos == 64 {
			c.generateBlock()
			c.keystreamPos = 0
		}

		take := 64 - c.keystreamPos
		if len(data)-offset < take {
			take = len(data) - offset
		}

		for i := 0; i < take; i++ {
			data[offset+i] ^= c.keystream[c.keystreamPos+i]
		}

		offset += take
		c.keystreamPos += take
	}
}

func (c *UxChaCha8) generateBlock() {
	var x [16]uint32
	copy(x[:], c.state[:])

	for i := 0; i < 4; i++ {
		quarterRound(&x[0], &x[4], &x[8], &x[12])
		quarterRound(&x[1], &x[5], &x[9], &x[13])
		quarterRound(&x[2], &x[6], &x[10], &x[14])
		quarterRound(&x[3], &x[7], &x[11], &x[15])

		quarterRound(&x[0], &x[5], &x[10], &x[15])
		quarterRound(&x[1], &x[6], &x[11], &x[12])
		quarterRound(&x[2], &x[7], &x[8], &x[13])
		quarterRound(&x[3], &x[4], &x[9], &x[14])
	}

	for i := 0; i < 16; i++ {
		binary.LittleEndian.PutUint32(c.keystream[i*4:(i+1)*4], x[i]+c.state[i])
	}

	c.state[12]++
}

func quarterRound(a, b, c, d *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)
}