#!/usr/bin/env python3
"""
Generate golden test vectors for AES-128 encryption testbench.
Uses pycryptodome library for reference implementation.
"""

from Crypto.Cipher import AES

# Test vectors - using FIPS-197 example values
TEST_CASES = [
    {
        "name": "FIPS-197 Appendix B",
        "key":       bytes.fromhex("2b7e151628aed2a6abf7158809cf4f3c"),
        "plaintext": bytes.fromhex("3243f6a8885a308d313198a2e0370734"),
        # Expected: 3925841d02dc09fbdc118597196a0b32
    },
    {
        "name": "All zeros",
        "key":       bytes.fromhex("00000000000000000000000000000000"),
        "plaintext": bytes.fromhex("00000000000000000000000000000000"),
        # Expected: 66e94bd4ef8a2c3b884cfa59ca342b2e
    },
    {
        "name": "All ones",
        "key":       bytes.fromhex("ffffffffffffffffffffffffffffffff"),
        "plaintext": bytes.fromhex("ffffffffffffffffffffffffffffffff"),
        # Expected: a1f6258c877d5fcd8964484538bfc92c
    },
]


def aes_encrypt(key: bytes, plaintext: bytes) -> bytes:
    """Encrypt using AES-128 ECB mode."""
    cipher = AES.new(key, AES.MODE_ECB)
    return cipher.encrypt(plaintext)


def bytes_to_verilog_words(data: bytes) -> list:
    """Convert 16 bytes to 4 x 32-bit words (big-endian)."""
    words = []
    for i in range(0, 16, 4):
        word = int.from_bytes(data[i:i+4], 'big')
        words.append(f"32'h{word:08X}")
    return words


def generate_verbose():
    """Generate verbose output with test vectors."""
    print("// ============================================")
    print("// AES-128 Encryption Test Vectors")
    print("// Generated using pycryptodome")
    print("// ============================================")
    print()
    
    for i, tc in enumerate(TEST_CASES):
        ciphertext = aes_encrypt(tc["key"], tc["plaintext"])
        
        print(f"// Test Case {i}: {tc['name']}")
        print(f"// Key:        {tc['key'].hex()}")
        print(f"// Plaintext:  {tc['plaintext'].hex()}")
        print(f"// Ciphertext: {ciphertext.hex()}")
        print()
        
        key_words = bytes_to_verilog_words(tc["key"])
        pt_words = bytes_to_verilog_words(tc["plaintext"])
        ct_words = bytes_to_verilog_words(ciphertext)
        
        print(f"// Verilog words (32-bit, big-endian):")
        print(f"// Key:        {{{', '.join(key_words)}}}")
        print(f"// Plaintext:  {{{', '.join(pt_words)}}}")
        print(f"// Ciphertext: {{{', '.join(ct_words)}}}")
        print()


def generate_testbench_data():
    """Generate Verilog array literals for testbench embedding."""
    print("    // Golden test vectors from generate_golden.py using pycryptodome")
    print("    // DO NOT EDIT - regenerate with: python3 generate_golden.py --array")
    print()
    
    # FIPS-197 test vector
    tc = TEST_CASES[0]
    ciphertext = aes_encrypt(tc["key"], tc["plaintext"])
    
    key_words = bytes_to_verilog_words(tc["key"])
    pt_words = bytes_to_verilog_words(tc["plaintext"])
    ct_words = bytes_to_verilog_words(ciphertext)
    
    print(f"    // Test: {tc['name']}")
    print(f"    // Key: {tc['key'].hex()}")
    print(f"    reg [31:0] test_key [0:3] = '{{{', '.join(key_words)}}};")
    print()
    print(f"    // Plaintext: {tc['plaintext'].hex()}")
    print(f"    reg [31:0] test_plaintext [0:3] = '{{{', '.join(pt_words)}}};")
    print()
    print(f"    // Expected ciphertext: {ciphertext.hex()}")
    print(f"    reg [31:0] expected_ciphertext [0:3] = '{{{', '.join(ct_words)}}};")


if __name__ == "__main__":
    import sys
    if len(sys.argv) > 1 and sys.argv[1] == "--array":
        generate_testbench_data()
    else:
        generate_verbose()