"""Tests for AES-256-CTR encryption, mirroring the Rust SDK test suite."""

from zg_storage.core.data import BytesDataSource
from zg_storage.core.encrypted_data import EncryptedDataSource
from zg_storage.transfer.encryption import (
    ENCRYPTION_HEADER_SIZE,
    ENCRYPTION_VERSION,
    EncryptionHeader,
    crypt_at,
    decrypt_file,
    decrypt_segment,
    encrypt_file,
)


def test_header_roundtrip():
    header = EncryptionHeader()
    raw = header.to_bytes()
    assert len(raw) == ENCRYPTION_HEADER_SIZE
    parsed = EncryptionHeader.parse(raw)
    assert parsed.version == ENCRYPTION_VERSION
    assert parsed.nonce == header.nonce


def test_crypt_roundtrip():
    key = bytes([0x42] * 32)
    nonce = bytes([0x13] * 16)
    original = b"hello world encryption test data"

    buf = bytearray(original)
    crypt_at(key, nonce, 0, buf)
    assert bytes(buf) != original

    # Decrypt (same operation for CTR)
    crypt_at(key, nonce, 0, buf)
    assert bytes(buf) == original


def test_crypt_at_offset():
    key = bytes([0x42] * 32)
    nonce = bytes([0x13] * 16)
    original = bytes(100)

    # Encrypt full
    full = bytearray(original)
    crypt_at(key, nonce, 0, full)

    # Encrypt in two parts at different offsets
    part1 = bytearray(original[:50])
    part2 = bytearray(original[50:])
    crypt_at(key, nonce, 0, part1)
    crypt_at(key, nonce, 50, part2)

    assert bytes(full[:50]) == bytes(part1)
    assert bytes(full[50:]) == bytes(part2)


def test_encrypt_decrypt_file():
    key = bytes([0x42] * 32)
    original = b"test data for encryption"

    encrypted = encrypt_file(key, original)
    assert len(encrypted) == len(original) + ENCRYPTION_HEADER_SIZE

    decrypted = decrypt_file(key, encrypted)
    assert decrypted == original


def test_decrypt_file_manual():
    """Build encrypted file manually and decrypt, matching Rust test."""
    key = bytes([0x42] * 32)
    original = b"test data for encryption"

    header = EncryptionHeader()
    encrypted_data = bytearray(original)
    crypt_at(key, header.nonce, 0, encrypted_data)

    encrypted_file = header.to_bytes() + bytes(encrypted_data)
    decrypted = decrypt_file(key, encrypted_file)
    assert decrypted == original


def test_decrypt_segment_first():
    """First segment includes the header."""
    key = bytes([0x42] * 32)
    original = b"segment zero data here"
    segment_size = 256 * 1024

    header = EncryptionHeader()
    encrypted_data = bytearray(original)
    crypt_at(key, header.nonce, 0, encrypted_data)

    segment_data = header.to_bytes() + bytes(encrypted_data)
    decrypted = decrypt_segment(key, 0, segment_size, segment_data, header)
    assert decrypted == original


def test_decrypt_segment_nonzero():
    """Non-zero segments have an offset-based data position."""
    key = bytes([0x42] * 32)
    segment_size = 256
    header = EncryptionHeader()

    # Simulate full data
    full_data = bytes(range(256)) * 4  # 1024 bytes
    full_encrypted = bytearray(full_data)
    crypt_at(key, header.nonce, 0, full_encrypted)

    # Segment 1 starts at offset segment_size - ENCRYPTION_HEADER_SIZE in the data stream
    data_offset = segment_size - ENCRYPTION_HEADER_SIZE
    segment_data = bytes(full_encrypted[data_offset : data_offset + segment_size])

    decrypted = decrypt_segment(key, 1, segment_size, segment_data, header)
    assert decrypted == full_data[data_offset : data_offset + segment_size]


def test_encrypted_data_source_size():
    original = bytes([1] * 1000)
    inner = BytesDataSource(original)
    key = bytes([0x42] * 32)
    encrypted = EncryptedDataSource(inner, key)
    assert encrypted.size() == len(original) + ENCRYPTION_HEADER_SIZE


def test_encrypted_data_source_read_header():
    original = bytes([1] * 100)
    inner = BytesDataSource(original)
    key = bytes([0x42] * 32)
    encrypted = EncryptedDataSource(inner, key)

    buf = encrypted.read_at(0, ENCRYPTION_HEADER_SIZE)
    assert len(buf) == ENCRYPTION_HEADER_SIZE
    assert buf[0] == ENCRYPTION_VERSION
    assert buf[1:17] == encrypted.header.nonce


def test_encrypted_data_source_roundtrip():
    original = b"hello world encryption test with EncryptedDataSource wrapper"
    inner = BytesDataSource(original)
    key = bytes([0x42] * 32)
    encrypted = EncryptedDataSource(inner, key)

    encrypted_size = encrypted.size()
    encrypted_buf = encrypted.read_at(0, encrypted_size)
    assert len(encrypted_buf) == encrypted_size

    decrypted = decrypt_file(key, encrypted_buf)
    assert decrypted == original


def test_encrypted_data_source_read_at_offset():
    original = bytes([0xAB] * 500)
    inner = BytesDataSource(original)
    key = bytes([0x42] * 32)
    encrypted = EncryptedDataSource(inner, key)

    encrypted_size = encrypted.size()
    full_buf = encrypted.read_at(0, encrypted_size)

    # Read in two parts and verify they match
    split = 100
    part1 = encrypted.read_at(0, split)
    part2 = encrypted.read_at(split, encrypted_size - split)

    assert full_buf[:split] == part1
    assert full_buf[split:] == part2
