// tb_mlkem_two_inst_xsim.v - ML-KEM hello_world on TWO mlkem_top instances.
//
// Models the real two-party split:
//   * Instance A (u_genenc): Alice/Bob side that does KeyGen THEN Encaps.
//       - KeyGen(d,z) writes ek into A's own ek_bram and produces dk.
//       - Encaps(m) reuses that same ek_bram (no re-streaming) -> shared_key, ct.
//   * Instance B (u_dec):    Alice side that does Decaps.
//       - receives dk (read out of A via dbg_dk) and ct (read out of A via
//         dbg_ct), streamed into B's input ports -> recovered_key.
//
// Then XOR-encrypt/decrypt "hello world" and verify both keys + message match.
// Each step prints its inputs and outputs.
//
//   xelab tb_mlkem_two_inst_xsim ; xsim
`timescale 1ns/1ps
module tb_mlkem_two_inst_xsim;
    localparam KP   = 2;                 // ML-KEM-512
    localparam EKB  = 384*KP + 32;       // 800
    localparam DKB  = 768*KP + 96;       // 1632
    localparam CTB  = 32*(10*KP + 4);    // 768
    localparam MLEN = 11;                // "hello world"

    reg clk=0;
    always #5 clk = ~clk;

    // ---------------- Instance A: KeyGen + Encaps ----------------
    reg         a_rst_n=0, a_start=0;
    reg [1:0]   a_op=0;
    reg [255:0] a_d=0, a_z=0, a_msg=0;
    wire        a_busy, a_done;
    wire [255:0] a_ss;
    reg  [10:0] a_ct_idx=0;  wire [7:0] a_ct_o;
    reg         a_byte_sel=0; reg [10:0] a_byte_idx=0; wire [7:0] a_byte_o;
    reg  [11:0] a_dk_idx=0;   wire [7:0] a_dk_o;
    wire [255:0] a_rho,a_sigma,a_r,a_hek,a_mp,a_kbar,a_dz,a_dh;

    mlkem_top u_genenc (
        .clk(clk), .rst_n(a_rst_n), .k_i(KP[2:0]), .op_i(a_op),
        .d_i(a_d), .z_i(a_z), .msg_i(a_msg), .start_i(a_start),
        .busy_o(a_busy), .done_o(a_done),
        .ek_in_we(1'b0), .ek_in_addr(11'd0), .ek_in_byte(8'd0),
        .dk_in_we(1'b0), .dk_in_addr(12'd0), .dk_in_byte(8'd0),
        .c_in_we(1'b0), .c_in_addr(11'd0), .c_in_byte(8'd0),
        .ss_o(a_ss), .dbg_ct_idx_i(a_ct_idx), .dbg_ct_o(a_ct_o),
        .dbg_slot_i(6'd0), .dbg_idx_i(8'd0), .dbg_coeff_o(),
        .dbg_byte_sel_i(a_byte_sel), .dbg_byte_idx_i(a_byte_idx), .dbg_byte_o(a_byte_o),
        .dbg_dk_idx_i(a_dk_idx), .dbg_dk_o(a_dk_o),
        .dbg_rho_o(a_rho), .dbg_sigma_o(a_sigma), .dbg_r_o(a_r), .dbg_hek_o(a_hek),
        .dbg_mprime_o(a_mp), .dbg_kbar_o(a_kbar), .dbg_decz_o(a_dz), .dbg_dech_o(a_dh)
    );

    // ---------------- Instance B: Decaps ----------------
    reg         b_rst_n=0, b_start=0;
    wire        b_busy, b_done;
    wire [255:0] b_ss;
    reg         b_dk_we=0; reg [11:0] b_dk_addr=0; reg [7:0] b_dk_byte=0;
    reg         b_c_we=0;  reg [10:0] b_c_addr=0;  reg [7:0] b_c_byte=0;
    wire [255:0] b_rho,b_sigma,b_r,b_hek,b_mp,b_kbar,b_dz,b_dh;
    wire [7:0]  b_byte_o, b_dk_o, b_ct_o;

    mlkem_top u_dec (
        .clk(clk), .rst_n(b_rst_n), .k_i(KP[2:0]), .op_i(2'd2),
        .d_i(256'd0), .z_i(256'd0), .msg_i(256'd0), .start_i(b_start),
        .busy_o(b_busy), .done_o(b_done),
        .ek_in_we(1'b0), .ek_in_addr(11'd0), .ek_in_byte(8'd0),
        .dk_in_we(b_dk_we), .dk_in_addr(b_dk_addr), .dk_in_byte(b_dk_byte),
        .c_in_we(b_c_we), .c_in_addr(b_c_addr), .c_in_byte(b_c_byte),
        .ss_o(b_ss), .dbg_ct_idx_i(11'd0), .dbg_ct_o(b_ct_o),
        .dbg_slot_i(6'd0), .dbg_idx_i(8'd0), .dbg_coeff_o(),
        .dbg_byte_sel_i(1'b0), .dbg_byte_idx_i(11'd0), .dbg_byte_o(b_byte_o),
        .dbg_dk_idx_i(12'd0), .dbg_dk_o(b_dk_o),
        .dbg_rho_o(b_rho), .dbg_sigma_o(b_sigma), .dbg_r_o(b_r), .dbg_hek_o(b_hek),
        .dbg_mprime_o(b_mp), .dbg_kbar_o(b_kbar), .dbg_decz_o(b_dz), .dbg_dech_o(b_dh)
    );

    // shuttle storage (the wire between the two instances)
    reg [7:0] dk_b  [0:DKB-1];
    reg [7:0] ct_b  [0:CTB-1];
    reg [7:0] msg_b [0:MLEN-1];
    reg [7:0] enc_b [0:MLEN-1];
    reg [7:0] dec_b [0:MLEN-1];
    reg [255:0] shared_key, recovered_key;
    integer c, i, j, errors;

    task wait_a; begin
        c=0; while(!a_done && c<4000000) begin @(posedge clk); c=c+1; end
        if(!a_done) begin $display("FAIL: instance A timeout (op=%0d)", a_op); $finish; end
    end endtask
    task wait_b; begin
        c=0; while(!b_done && c<4000000) begin @(posedge clk); c=c+1; end
        if(!b_done) begin $display("FAIL: instance B timeout"); $finish; end
    end endtask

    task print_hex32; input [255:0] v; begin
        for (j=0;j<32;j=j+1) $write("%02x", v[8*j +: 8]); $write("\n");
    end endtask

    initial begin
        errors = 0;
        msg_b[0]="h"; msg_b[1]="e"; msg_b[2]="l"; msg_b[3]="l"; msg_b[4]="o";
        msg_b[5]=" "; msg_b[6]="w"; msg_b[7]="o"; msg_b[8]="r"; msg_b[9]="l"; msg_b[10]="d";

        $display("=== ML-KEM hello_world (TWO instances: genenc + dec, ML-KEM-512) ===");
        $write("Original: \""); for(i=0;i<MLEN;i=i+1) $write("%c", msg_b[i]); $display("\"\n");

        a_rst_n=0; b_rst_n=0; repeat(4) @(posedge clk); a_rst_n=1; b_rst_n=1; @(posedge clk);

        // ===== Step 1: Instance A — KeyGen(d, z) =====
        a_d = {32{8'h42}}; a_z = {32{8'h77}};
        $display("[A.1] KeyGen");
        $write("      in  d = "); print_hex32(a_d);
        $write("      in  z = "); print_hex32(a_z);
        a_op=2'd0; a_start=1; @(posedge clk); a_start=0; wait_a;
        // read dk out of A (to hand to instance B)
        for (i=0;i<DKB;i=i+1) begin a_dk_idx=i[11:0]; @(posedge clk); @(posedge clk); dk_b[i]=a_dk_o; end
        $display("      out ek = %0d B (stays in A's ek_bram); dk = %0d B (dk[0:8]=%02x %02x %02x %02x %02x %02x %02x %02x ...)",
                 EKB, DKB, dk_b[0],dk_b[1],dk_b[2],dk_b[3],dk_b[4],dk_b[5],dk_b[6],dk_b[7]);
        $display("      (cycles: %0d)\n", c);

        // ===== Step 2: Instance A — Encaps(ek, m) =====
        // ek already lives in A's ek_bram from KeyGen; no re-streaming needed.
        a_msg = {32{8'hDE}};
        $display("[A.2] Encaps (reuses ek_bram from KeyGen)");
        $write("      in  m = "); print_hex32(a_msg);
        a_op=2'd1; a_start=1; @(posedge clk); a_start=0; wait_a;
        shared_key = a_ss;
        for (i=0;i<CTB;i=i+1) begin a_ct_idx=i[10:0]; @(posedge clk); @(posedge clk); ct_b[i]=a_ct_o; end
        $write("      out shared_key = "); print_hex32(shared_key);
        $display("      out kem_ct = %0d B (ct[0:8]=%02x %02x %02x %02x %02x %02x %02x %02x ...)",
                 CTB, ct_b[0],ct_b[1],ct_b[2],ct_b[3],ct_b[4],ct_b[5],ct_b[6],ct_b[7]);
        $display("      (cycles: %0d)\n", c);

        // ===== Step 3: Bob — XOR-encrypt =====
        for (i=0;i<MLEN;i=i+1) enc_b[i] = msg_b[i] ^ shared_key[8*(i%32) +: 8];
        $display("[3] encrypt \"hello world\" (XOR stream, demo)");
        $write("    out encrypted = "); for(i=0;i<MLEN;i=i+1) $write("%02x ", enc_b[i]); $write("\n\n");

        // ===== Step 4: Instance B — Decaps(dk, kem_ct) =====
        // hand dk + ct from A to B over the streaming input ports
        for (i=0;i<DKB;i=i+1) begin b_dk_we=1; b_dk_addr=i[11:0]; b_dk_byte=dk_b[i]; @(posedge clk); end
        b_dk_we=0;
        for (i=0;i<CTB;i=i+1) begin b_c_we=1; b_c_addr=i[10:0]; b_c_byte=ct_b[i]; @(posedge clk); end
        b_c_we=0; @(posedge clk);
        $display("[B.1] Decaps (separate instance)");
        $display("      in  dk = %0d B, kem_ct = %0d B  (received from instance A)", DKB, CTB);
        b_start=1; @(posedge clk); b_start=0; wait_b;
        recovered_key = b_ss;
        $write("      out recovered_key = "); print_hex32(recovered_key);
        $display("      (cycles: %0d)\n", c);

        // ===== Step 5: Alice — XOR-decrypt =====
        for (i=0;i<MLEN;i=i+1) dec_b[i] = enc_b[i] ^ recovered_key[8*(i%32) +: 8];
        $write("[5] decrypt -> \""); for(i=0;i<MLEN;i=i+1) $write("%c", dec_b[i]); $display("\"\n");

        // ===== Step 6: Verify =====
        if (shared_key !== recovered_key) begin
            $display("FAIL: shared_key (A) != recovered_key (B)"); errors=errors+1;
        end
        for (i=0;i<MLEN;i=i+1) if (dec_b[i] !== msg_b[i]) errors=errors+1;

        if (errors==0)
            $display("Success: A.shared_key == B.recovered_key, message recovered across two instances.");
        else
            $display("FAILURE: %0d mismatches", errors);
        $finish;
    end

    initial begin #400000000; $display("FAIL: global timeout"); $finish; end
endmodule