"""Tests for i2p_crypto.aes — AES-256-CBC engine.

Covers:
  - Round-trip encrypt/decrypt (CBC)
  - Round-trip encrypt_block/decrypt_block (single block ECB)
  - Known NIST test vectors (AES-256-CBC)
  - Wrong-key detection
  - Block-alignment enforcement (payload not a multiple of 16)
  - IV length enforcement
  - Empty payload rejection
"""

from __future__ import annotations

import os

import pytest

from i2p_crypto.aes import AESEngine


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _random_key() -> bytes:
    """Return a random 32-byte AES-256 key."""
    return os.urandom(32)


def _random_iv() -> bytes:
    """Return a random 16-byte IV."""
    return os.urandom(16)


# ---------------------------------------------------------------------------
# Round-trip tests
# ---------------------------------------------------------------------------

class TestRoundTrip:
    """encrypt then decrypt must return the original plaintext."""

    def test_single_block(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(16)
        ciphertext = AESEngine.encrypt(plaintext, key, iv)
        assert AESEngine.decrypt(ciphertext, key, iv) == plaintext

    def test_multi_block(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(16 * 8)  # 128 bytes
        ciphertext = AESEngine.encrypt(plaintext, key, iv)
        assert len(ciphertext) == len(plaintext)
        assert AESEngine.decrypt(ciphertext, key, iv) == plaintext

    def test_large_payload(self):
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(16 * 64)  # 1024 bytes
        ciphertext = AESEngine.encrypt(plaintext, key, iv)
        assert AESEngine.decrypt(ciphertext, key, iv) == plaintext

    def test_in_place_round_trip(self):
        """Mirrors Java's testED2: encrypt, then decrypt, verify equality."""
        key, iv = _random_key(), _random_iv()
        orig = os.urandom(128)
        data = AESEngine.encrypt(orig, key, iv)
        data = AESEngine.decrypt(data, key, iv)
        assert data == orig

    def test_deterministic(self):
        """Same key+IV+plaintext must produce same ciphertext."""
        key, iv = _random_key(), _random_iv()
        plaintext = os.urandom(48)
        c1 = AESEngine.encrypt(plaintext, key, iv)
        c2 = AESEngine.encrypt(plaintext, key, iv)
        assert c1 == c2


# ---------------------------------------------------------------------------
# Single-block (ECB) tests
# ---------------------------------------------------------------------------

class TestBlockOperations:
    """encrypt_block / decrypt_block round-trip and validation."""

    def test_round_trip(self):
        key = _random_key()
        block = os.urandom(16)
        enc = AESEngine.encrypt_block(block, key)
        assert AESEngine.decrypt_block(enc, key) == block

    def test_wrong_size_encrypt(self):
        key = _random_key()
        with pytest.raises(ValueError, match="exactly 16 bytes"):
            AESEngine.encrypt_block(os.urandom(15), key)

    def test_wrong_size_decrypt(self):
        key = _random_key()
        with pytest.raises(ValueError, match="exactly 16 bytes"):
            AESEngine.decrypt_block(os.urandom(17), key)


# ---------------------------------------------------------------------------
# Known test vectors — NIST AES-256-CBC (F.2.5 / F.2.6)
# ---------------------------------------------------------------------------

class TestKnownVectors:
    """NIST SP 800-38A AES-256-CBC test vectors."""

    # Key, IV, plaintext blocks, expected ciphertext blocks from NIST SP 800-38A
    NIST_KEY = bytes.fromhex(
        "603deb1015ca71be2b73aef0857d7781"
        "1f352c073b6108d72d9810a30914dff4"
    )
    NIST_IV = bytes.fromhex("000102030405060708090a0b0c0d0e0f")

    NIST_PLAINTEXT = bytes.fromhex(
        "6bc1bee22e409f96e93d7e117393172a"
        "ae2d8a571e03ac9c9eb76fac45af8e51"
        "30c81c46a35ce411e5fbc1191a0a52ef"
        "f69f2445df4f9b17ad2b417be66c3710"
    )

    NIST_CIPHERTEXT = bytes.fromhex(
        "f58c4c04d6e5f1ba779eabfb5f7bfbd6"
        "9cfc4e967edb808d679f777bc6702c7d"
        "39f23369a9d9bacfa530e26304231461"
        "b2eb05e2c39be9fcda6c19078c6a9d1b"
    )

    def test_encrypt(self):
        ct = AESEngine.encrypt(self.NIST_PLAINTEXT, self.NIST_KEY, self.NIST_IV)
        assert ct == self.NIST_CIPHERTEXT

    def test_decrypt(self):
        pt = AESEngine.decrypt(self.NIST_CIPHERTEXT, self.NIST_KEY, self.NIST_IV)
        assert pt == self.NIST_PLAINTEXT


# ---------------------------------------------------------------------------
# Wrong-key detection
# ---------------------------------------------------------------------------

class TestWrongKey:
    """Decrypting with a different key must not return the original plaintext."""

    def test_wrong_key_cbc(self):
        key1, key2 = _random_key(), _random_key()
        iv = _random_iv()
        plaintext = os.urandom(64)
        ciphertext = AESEngine.encrypt(plaintext, key1, iv)
        decrypted = AESEngine.decrypt(ciphertext, key2, iv)
        assert decrypted != plaintext

    def test_wrong_key_block(self):
        key1, key2 = _random_key(), _random_key()
        block = os.urandom(16)
        enc = AESEngine.encrypt_block(block, key1)
        dec = AESEngine.decrypt_block(enc, key2)
        assert dec != block


# ---------------------------------------------------------------------------
# Alignment / validation enforcement
# ---------------------------------------------------------------------------

class TestValidation:
    """Payload length and IV size must be enforced."""

    def test_non_aligned_encrypt(self):
        with pytest.raises(ValueError, match="multiple of 16"):
            AESEngine.encrypt(os.urandom(15), _random_key(), _random_iv())

    def test_non_aligned_decrypt(self):
        with pytest.raises(ValueError, match="multiple of 16"):
            AESEngine.decrypt(os.urandom(17), _random_key(), _random_iv())

    def test_empty_payload_encrypt(self):
        with pytest.raises(ValueError, match="empty"):
            AESEngine.encrypt(b"", _random_key(), _random_iv())

    def test_empty_payload_decrypt(self):
        with pytest.raises(ValueError, match="empty"):
            AESEngine.decrypt(b"", _random_key(), _random_iv())

    def test_bad_iv_length(self):
        with pytest.raises(ValueError, match="IV must be exactly 16"):
            AESEngine.encrypt(os.urandom(16), _random_key(), os.urandom(12))

    def test_odd_size_payload(self):
        """Payload of 7 bytes — not aligned."""
        with pytest.raises(ValueError, match="multiple of 16"):
            AESEngine.encrypt(os.urandom(7), _random_key(), _random_iv())

    def test_33_byte_payload(self):
        """33 bytes — one byte over two blocks."""
        with pytest.raises(ValueError, match="multiple of 16"):
            AESEngine.encrypt(os.urandom(33), _random_key(), _random_iv())


# ---------------------------------------------------------------------------
# Key size variants
# ---------------------------------------------------------------------------

class TestKeySizes:
    """AES supports 128, 192, and 256-bit keys."""

    @pytest.mark.parametrize("key_len", [16, 24, 32])
    def test_round_trip_key_sizes(self, key_len: int):
        key = os.urandom(key_len)
        iv = _random_iv()
        plaintext = os.urandom(48)
        ct = AESEngine.encrypt(plaintext, key, iv)
        assert AESEngine.decrypt(ct, key, iv) == plaintext