package crypto

import (
	"bytes"
	"testing"
)

func TestCTRFromSecretRoundTrip(t *testing.T) {
	secret := Hash([]byte("test-secret-key"))
	nonce := []byte("test-nonce-001")
	plaintext := []byte("Hello, hamadryad CTR encryption!")

	stream := NewCTRStreamFromSecret(secret, nonce)
	ciphertext := stream.Encrypt(plaintext)

	if len(ciphertext) != len(plaintext) {
		t.Errorf("ciphertext length = %d, want %d", len(ciphertext), len(plaintext))
	}
	if bytes.Equal(ciphertext, plaintext) {
		t.Error("ciphertext should differ from plaintext")
	}

	stream2 := NewCTRStreamFromSecret(secret, nonce)
	recovered := stream2.Decrypt(ciphertext)
	if !bytes.Equal(recovered, plaintext) {
		t.Errorf("recovered = %q, want %q", recovered, plaintext)
	}
}

func TestCTRDeterministic(t *testing.T) {
	secret := Hash([]byte("determinism-key"))
	nonce := []byte("determinism-nonce")
	plaintext := []byte("same message twice")

	s1 := NewCTRStreamFromSecret(secret, nonce)
	s2 := NewCTRStreamFromSecret(secret, nonce)

	ct1 := s1.Encrypt(plaintext)
	ct2 := s2.Encrypt(plaintext)
	if !bytes.Equal(ct1, ct2) {
		t.Error("CTR encryption should be deterministic with same key+nonce")
	}
}

func TestCTRDifferentNonces(t *testing.T) {
	secret := Hash([]byte("nonce-test-key"))
	plaintext := []byte("same plaintext")

	s1 := NewCTRStreamFromSecret(secret, []byte("nonce-1"))
	s2 := NewCTRStreamFromSecret(secret, []byte("nonce-2"))

	ct1 := s1.Encrypt(plaintext)
	ct2 := s2.Encrypt(plaintext)
	if bytes.Equal(ct1, ct2) {
		t.Error("different nonces should produce different ciphertext")
	}
}

func TestCTRDifferentSecrets(t *testing.T) {
	secret1 := Hash([]byte("secret-1"))
	secret2 := Hash([]byte("secret-2"))
	nonce := []byte("same-nonce")
	plaintext := []byte("same plaintext")

	s1 := NewCTRStreamFromSecret(secret1, nonce)
	s2 := NewCTRStreamFromSecret(secret2, nonce)

	ct1 := s1.Encrypt(plaintext)
	ct2 := s2.Encrypt(plaintext)
	if bytes.Equal(ct1, ct2) {
		t.Error("different secrets should produce different ciphertext")
	}

	// Cross-key decryption should fail.
	s3 := NewCTRStreamFromSecret(secret2, nonce)
	wrong := s3.Decrypt(ct1)
	if bytes.Equal(wrong, plaintext) {
		t.Error("wrong secret should not decrypt correctly")
	}
}

func TestCTRRandomAccess(t *testing.T) {
	secret := Hash([]byte("random-access-key"))
	stream := NewCTRStreamFromSecret(secret, []byte("random-access"))

	block0 := stream.KeystreamBlock(0)
	block1 := stream.KeystreamBlock(1)
	block2 := stream.KeystreamBlock(2)

	// Order-independent.
	block2r := stream.KeystreamBlock(2)
	block1r := stream.KeystreamBlock(1)
	block0r := stream.KeystreamBlock(0)

	if !bytes.Equal(block0, block0r) {
		t.Error("block 0 should be same regardless of generation order")
	}
	if !bytes.Equal(block1, block1r) {
		t.Error("block 1 should be same regardless of generation order")
	}
	if !bytes.Equal(block2, block2r) {
		t.Error("block 2 should be same regardless of generation order")
	}
	if bytes.Equal(block0, block1) {
		t.Error("block 0 and 1 should differ")
	}
	if bytes.Equal(block1, block2) {
		t.Error("block 1 and 2 should differ")
	}
}

func TestCTRRandomAccessEncrypt(t *testing.T) {
	secret := Hash([]byte("seek-test-key"))
	nonce := []byte("seek-test")

	plaintext := make([]byte, BlockSize*3)
	for i := range plaintext {
		plaintext[i] = byte(i % 256)
	}

	stream := NewCTRStreamFromSecret(secret, nonce)
	ciphertext := stream.Encrypt(plaintext)

	stream2 := NewCTRStreamFromSecret(secret, nonce)
	blockSlice := ciphertext[BlockSize : BlockSize*2]
	decrypted := stream2.EncryptCTRAt(blockSlice, uint64(BlockSize))

	expected := plaintext[BlockSize : BlockSize*2]
	if !bytes.Equal(decrypted, expected) {
		t.Error("random-access decryption should recover correct block")
	}
}

func TestCTREmptyPlaintext(t *testing.T) {
	secret := Hash([]byte("empty-test"))
	stream := NewCTRStreamFromSecret(secret, []byte("nonce"))
	ct := stream.Encrypt([]byte{})
	if len(ct) != 0 {
		t.Errorf("ciphertext of empty plaintext should be empty, got %d bytes", len(ct))
	}
}

func TestCTRBlockSizeIs64(t *testing.T) {
	if BlockSize != 64 {
		t.Errorf("BlockSize = %d, want 64", BlockSize)
	}
}

func TestCTRKeystreamNonZero(t *testing.T) {
	secret := Hash([]byte("nonzero-key"))
	stream := NewCTRStreamFromSecret(secret, []byte("nonce"))
	block := stream.KeystreamBlock(0)

	allZero := true
	for _, b := range block {
		if b != 0 {
			allZero = false
			break
		}
	}
	if allZero {
		t.Error("keystream block should not be all zeros")
	}
}