`timescale 1ns / 1ps

// AES-128 Encryption Testbench
// Uses FIPS-197 Appendix B test vector
// Golden values from: python3 generate_golden.py --array

module tb_aes128_encrypt;
    reg clk, rst;
    reg load_key, load_data, start, read_data;
    reg [31:0] data_in;
    wire [31:0] data_out;
    wire done, busy;
    
    integer i, errors;
    parameter TIMEOUT = 500;
    
    // Golden test vectors from generate_golden.py using pycryptodome
    // FIPS-197 Appendix B test vector
    reg [31:0] test_key [0:3];
    reg [31:0] test_plaintext [0:3];
    reg [31:0] expected_ciphertext [0:3];

    initial begin
        test_key[0] = 32'h2B7E1516; test_key[1] = 32'h28AED2A6; test_key[2] = 32'hABF71588; test_key[3] = 32'h09CF4F3C;
        test_plaintext[0] = 32'h3243F6A8; test_plaintext[1] = 32'h885A308D; test_plaintext[2] = 32'h313198A2; test_plaintext[3] = 32'hE0370734;
        expected_ciphertext[0] = 32'h3925841D; expected_ciphertext[1] = 32'h02DC09FB; expected_ciphertext[2] = 32'hDC118597; expected_ciphertext[3] = 32'h196A0B32;
    end
    
    reg [31:0] result [0:3];
    
    // Instantiate UUT
    aes128_encrypt uut (
        .clk(clk), .rst(rst),
        .load_key(load_key), .load_data(load_data),
        .data_in(data_in), .start(start),
        .data_out(data_out), .read_data(read_data),
        .done(done), .busy(busy)
    );
    
    // Clock 100MHz
    initial begin clk = 0; forever #5 clk = ~clk; end
    
    // Test procedure
    initial begin
        errors = 0;
        rst = 1; load_key = 0; load_data = 0; start = 0; read_data = 0; data_in = 0;
        $display("=== AES-128 Encryption Testbench ===");
        $display("Key:       2b7e151628aed2a6abf7158809cf4f3c");
        $display("Plaintext: 3243f6a8885a308d313198a2e0370734");
        $display("Expected:  3925841d02dc09fbdc118597196a0b32");
        
        repeat(5) @(posedge clk);
        rst = 0;
        repeat(2) @(posedge clk);
        
        // Load key (4 x 32-bit words)
        $display("Loading key...");
        for (i = 0; i < 4; i = i + 1) begin
            @(posedge clk);
            load_key = 1;
            data_in = test_key[i];
            @(posedge clk);
            load_key = 0;
        end
        
        // Load plaintext (4 x 32-bit words)
        $display("Loading plaintext...");
        for (i = 0; i < 4; i = i + 1) begin
            @(posedge clk);
            load_data = 1;
            data_in = test_plaintext[i];
            @(posedge clk);
            load_data = 0;
        end
        
        // Start encryption
        $display("Starting encryption...");
        @(posedge clk);
        start = 1;
        @(posedge clk);
        start = 0;
        
        // Wait for done
        i = 0;
        while (!done && i < TIMEOUT) begin
            @(posedge clk);
            i = i + 1;
        end
        
        if (i >= TIMEOUT) begin
            $display("ERROR: Timeout waiting for done signal");
            errors = 1;
        end else begin
            $display("Encryption complete in %0d cycles", i);
        end
        
        // Read ciphertext (4 x 32-bit words)
        $display("Reading ciphertext...");
        for (i = 0; i < 4; i = i + 1) begin
            @(posedge clk);
            read_data = 1;
            @(posedge clk);
            result[i] = data_out;
            read_data = 0;
        end
        
        // Verify
        $display("Verifying results...");
        for (i = 0; i < 4; i = i + 1) begin
            if (result[i] !== expected_ciphertext[i]) begin
                $display("ERROR: Word %0d: Expected 0x%08X, got 0x%08X", 
                         i, expected_ciphertext[i], result[i]);
                errors = errors + 1;
            end else begin
                $display("  Word %0d: 0x%08X OK", i, result[i]);
            end
        end
        
        $display("");
        if (errors == 0)
            $display("TEST_RESULT: PASS");
        else
            $display("TEST_RESULT: FAIL (%0d errors)", errors);
        
        $finish;
    end

endmodule