package crypto

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

// CTR implements counter mode encryption using ChaCha20.
//
// A shared secret (Hamadryad) is hashed once to derive a 32-byte ChaCha20
// key. This is a one-time key-derivation step using SWIFFT's proven
// collision/preimage resistance -- the hash is NOT used as a PRF. All
// keystream generation is performed by ChaCha20.
//
// Random access: ChaCha20's SetCounter enables seeking in encrypted streams.

// BlockSize is the number of keystream bytes produced per ChaCha20 block counter.
const BlockSize = 64

// CTRStream holds the state for ChaCha20 stream encryption/decryption.
type CTRStream struct {
	chachaKey   [chacha20.KeySize]byte
	chachaNonce [chacha20.NonceSize]byte
}

// newCipher creates a fresh ChaCha20 cipher from the stream's key/nonce.
func (s *CTRStream) newCipher() *chacha20.Cipher {
	c, _ := chacha20.NewUnauthenticatedCipher(s.chachaKey[:], s.chachaNonce[:])
	return c
}

// deriveChaCha20Key derives a 32-byte ChaCha20 key from Hamadryad key material.
func deriveChaCha20Key(keyMaterial Hamadryad) [chacha20.KeySize]byte {
	keyHash := Hash(append([]byte("hamadryad-chacha20-key-v1"), keyMaterial[:]...))
	var key [chacha20.KeySize]byte
	copy(key[:], keyHash[:chacha20.KeySize])
	return key
}

// deriveChaCha20Nonce derives a 12-byte ChaCha20 nonce from Hamadryad nonce material.
func deriveChaCha20Nonce(nonceMaterial Hamadryad) [chacha20.NonceSize]byte {
	nonceHash := Hash(append([]byte("hamadryad-chacha20-nonce-v1"), nonceMaterial[:]...))
	var nonce [chacha20.NonceSize]byte
	copy(nonce[:], nonceHash[:chacha20.NonceSize])
	return nonce
}

// NewCTRStreamFromSecret creates a CTR stream using a shared secret as the key.
// The nonce should be unique per message for a given secret.
func NewCTRStreamFromSecret(secret Hamadryad, nonce []byte) *CTRStream {
	nonceHash := Hash(nonce)
	return &CTRStream{
		chachaKey:   deriveChaCha20Key(secret),
		chachaNonce: deriveChaCha20Nonce(nonceHash),
	}
}

// Encrypt encrypts plaintext using the ChaCha20 keystream.
func (s *CTRStream) Encrypt(plaintext []byte) []byte {
	return s.xor(plaintext)
}

// Decrypt decrypts ciphertext using the ChaCha20 keystream.
// Decryption is identical to encryption (XOR is self-inverse).
func (s *CTRStream) Decrypt(ciphertext []byte) []byte {
	return s.xor(ciphertext)
}

// KeystreamBlock generates the keystream block for a given counter.
// Each call is independent -- blocks can be generated in any order.
func (s *CTRStream) KeystreamBlock(counter uint64) []byte {
	c := s.newCipher()
	c.SetCounter(uint32(counter))
	block := make([]byte, BlockSize)
	c.XORKeyStream(block, block) // XOR zeros = raw keystream
	return block
}

// EncryptCTRAt encrypts a single block at the given byte offset.
// This enables random-access encryption of large streams.
func (s *CTRStream) EncryptCTRAt(plaintext []byte, offset uint64) []byte {
	c := s.newCipher()
	blockIdx := uint32(offset / uint64(BlockSize))
	blockOff := int(offset % uint64(BlockSize))
	c.SetCounter(blockIdx)

	// Discard partial block prefix to reach the target offset.
	if blockOff > 0 {
		discard := make([]byte, blockOff)
		c.XORKeyStream(discard, discard)
	}

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

// xor XORs data with the ChaCha20 keystream.
func (s *CTRStream) xor(data []byte) []byte {
	c := s.newCipher()
	result := make([]byte, len(data))
	c.XORKeyStream(result, data)
	return result
}