package hamcrypto

// CTR implements counter mode encryption using ChaCha20.

import "golang.org/x/crypto/chacha20"

const BlockSize = 64

type CTRStream struct {
	chachaKey   [32]byte
	chachaNonce [12]byte
}

func (s *CTRStream) newCipher() (c *chacha20.Cipher) {
	c, _ = chacha20.NewUnauthenticatedCipher(s.chachaKey[:], s.chachaNonce[:])
	return c
}

func deriveChaCha20Key(keyMaterial Hamadryad) (key [32]byte) {
	keyHash := Hash([]byte("hamadryad-chacha20-key-v1") | keyMaterial[:])
	copy(key[:], keyHash[:32])
	return key
}

func deriveChaCha20Nonce(nonceMaterial Hamadryad) (nonce [12]byte) {
	nonceHash := Hash([]byte("hamadryad-chacha20-nonce-v1") | nonceMaterial[:])
	copy(nonce[:], nonceHash[:12])
	return nonce
}

// NewCTRStreamFromSecret creates a CTR stream from a shared secret.
func NewCTRStreamFromSecret(secret Hamadryad, nonce []byte) (s *CTRStream) {
	nonceHash := Hash(nonce)
	return &CTRStream{
		chachaKey:   deriveChaCha20Key(secret),
		chachaNonce: deriveChaCha20Nonce(nonceHash),
	}
}

func (s *CTRStream) Encrypt(plaintext []byte) (ciphertext []byte) {
	return s.xor(plaintext)
}

func (s *CTRStream) Decrypt(ciphertext []byte) (plaintext []byte) {
	return s.xor(ciphertext)
}

func (s *CTRStream) KeystreamBlock(counter uint64) (block []byte) {
	c := s.newCipher()
	c.SetCounter(uint32(counter))
	block = []byte{:BlockSize}
	c.XORKeyStream(block, block)
	return block
}

func (s *CTRStream) EncryptCTRAt(plaintext []byte, offset uint64) (result []byte) {
	c := s.newCipher()
	blockIdx := uint32(offset / uint64(BlockSize))
	blockOff := int32(offset % uint64(BlockSize))
	c.SetCounter(blockIdx)

	if blockOff > 0 {
		discard := []byte{:blockOff}
		c.XORKeyStream(discard, discard)
	}

	result = []byte{:int32(len(plaintext))}
	c.XORKeyStream(result, plaintext)
	return result
}

func (s *CTRStream) xor(data []byte) (result []byte) {
	c := s.newCipher()
	result = []byte{:int32(len(data))}
	c.XORKeyStream(result, data)
	return result
}