`timescale 1ns / 1ps

// Pipelined AES-128 Encryption Testbench
// Tests II=1 pipelining with multiple blocks
// Golden values from pycryptodome

module tb_aes128_pipelined;
    reg clk, rst;
    reg [127:0] key;
    reg [127:0] plaintext;
    reg valid_in;
    wire [127:0] ciphertext;
    wire valid_out;
    
    integer i, errors, received;
    parameter NUM_TESTS = 4;
    parameter TIMEOUT = 100;
    
    // FIPS-197 Appendix B test key
    reg [127:0] test_key = 128'h2b7e151628aed2a6abf7158809cf4f3c;
    
    // Test plaintexts and expected ciphertexts
    reg [127:0] test_pt [0:3];
    reg [127:0] expected_ct [0:3];
    reg [127:0] result_ct [0:3];
    
    // Instantiate UUT
    aes128_pipelined uut (
        .clk(clk), .rst(rst),
        .key(key), .plaintext(plaintext), .valid_in(valid_in),
        .ciphertext(ciphertext), .valid_out(valid_out)
    );
    
    // Clock 100MHz
    initial begin clk = 0; forever #5 clk = ~clk; end
    
    initial begin
        // Test vectors (all using same key)
        // PT 0: FIPS-197 test
        test_pt[0] = 128'h3243f6a8885a308d313198a2e0370734;
        expected_ct[0] = 128'h3925841d02dc09fbdc118597196a0b32;
        // PT 1: All zeros
        test_pt[1] = 128'h00000000000000000000000000000000;
        expected_ct[1] = 128'h7df76b0c1ab899b33e42f047b91b546f;
        // PT 2: All ones
        test_pt[2] = 128'hffffffffffffffffffffffffffffffff;
        expected_ct[2] = 128'h8af2860142f786f409307c1a3f7eaaac;
        // PT 3: Sequential
        test_pt[3] = 128'h00112233445566778899aabbccddeeff;
        expected_ct[3] = 128'h8df4e9aac5c7573a27d8d055d6e4d64b;
    end
    
    // Test procedure
    initial begin
        errors = 0; received = 0;
        rst = 1; key = 0; plaintext = 0; valid_in = 0;
        $display("=== Pipelined AES-128 Encryption Testbench ===");
        $display("Testing II=1 pipelining with %0d blocks", NUM_TESTS);
        
        repeat(5) @(posedge clk);
        rst = 0;
        @(posedge clk);
        
        // Set key
        key = test_key;
        
        // Stream in all plaintexts (one per cycle for II=1)
        $display("Streaming %0d plaintext blocks...", NUM_TESTS);
        for (i = 0; i < NUM_TESTS; i = i + 1) begin
            @(posedge clk);
            plaintext = test_pt[i];
            valid_in = 1;
        end
        @(posedge clk);
        valid_in = 0;
        
        // Wait for outputs
        $display("Waiting for ciphertext outputs...");
        i = 0;
        while (received < NUM_TESTS && i < TIMEOUT) begin
            @(posedge clk);
            if (valid_out) begin
                result_ct[received] = ciphertext;
                received = received + 1;
            end
            i = i + 1;
        end
        
        if (received < NUM_TESTS) begin
            $display("ERROR: Timeout - received %0d/%0d blocks", received, NUM_TESTS);
            errors = NUM_TESTS - received;
        end
        
        // Verify results
        $display("Verifying %0d results...", received);
        for (i = 0; i < received; i = i + 1) begin
            if (result_ct[i] !== expected_ct[i]) begin
                $display("ERROR: Block %0d: Expected 0x%032h, got 0x%032h", 
                         i, expected_ct[i], result_ct[i]);
                errors = errors + 1;
            end else begin
                $display("  Block %0d: OK", i);
            end
        end
        
        $display("");
        if (errors == 0)
            $display("TEST_RESULT: PASS");
        else
            $display("TEST_RESULT: FAIL (%0d errors)", errors);
        
        $finish;
    end

endmodule