// tb_mlkem_hello_world_xsim.v - Hardware run of ml-kem-r examples/hello_world.rs.
//
// Mirrors the Rust example end-to-end on the mlkem_top DUT (ML-KEM-512, K=2):
//   1. Alice: KeyGen(d=0x42..., z=0x77...)            -> (ek 800B, dk 1632B)
//   2. Bob:   Encaps(ek, m=0xDE...)                   -> (shared_key, kem_ct 768B)
//   3. Bob:   XOR-encrypt "hello world" with key      -> encrypted
//   4. Alice: Decaps(dk, kem_ct)                       -> recovered_key
//   5. Alice: XOR-decrypt encrypted with key          -> "hello world"
//   6. Verify: shared_key == recovered_key, decrypted == message
//
// The whole protocol runs on ONE DUT instance: ek/dk are read out of KeyGen via
// the dbg taps, then fed back into Encaps/Decaps via the streaming input ports,
// exactly as a real ek/dk/ct would cross the wire between Alice and Bob.
//
// Each step prints its inputs and outputs.
//
//   xelab tb_mlkem_hello_world_xsim ; xsim
`timescale 1ns/1ps
module tb_mlkem_hello_world_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, rst_n=0, start_i=0;
    reg [2:0]   k_i = KP;
    reg [1:0]   op_i = 0;
    reg [255:0] d_i=0, z_i=0, msg_i=0;
    wire busy_o, done_o;
    reg         ek_in_we=0; reg [10:0] ek_in_addr=0; reg [7:0] ek_in_byte=0;
    reg         dk_in_we=0; reg [11:0] dk_in_addr=0; reg [7:0] dk_in_byte=0;
    reg         c_in_we=0;  reg [10:0] c_in_addr=0;  reg [7:0] c_in_byte=0;
    wire [255:0] ss_o;
    reg  [10:0] dbg_ct_idx_i=0; wire [7:0] dbg_ct_o;
    reg [5:0] dbg_slot_i=0; reg [7:0] dbg_idx_i=0; wire [11:0] dbg_coeff_o;
    reg dbg_byte_sel_i=0; reg [10:0] dbg_byte_idx_i=0; wire [7:0] dbg_byte_o;
    reg [11:0] dbg_dk_idx_i=0; wire [7:0] dbg_dk_o;
    wire [255:0] dbg_rho_o, dbg_sigma_o, dbg_r_o, dbg_hek_o;
    wire [255:0] dbg_mprime_o, dbg_kbar_o, dbg_decz_o, dbg_dech_o;

    mlkem_top dut (
        .clk(clk), .rst_n(rst_n), .k_i(k_i), .op_i(op_i),
        .d_i(d_i), .z_i(z_i), .msg_i(msg_i), .start_i(start_i),
        .busy_o(busy_o), .done_o(done_o),
        .ek_in_we(ek_in_we), .ek_in_addr(ek_in_addr), .ek_in_byte(ek_in_byte),
        .dk_in_we(dk_in_we), .dk_in_addr(dk_in_addr), .dk_in_byte(dk_in_byte),
        .c_in_we(c_in_we), .c_in_addr(c_in_addr), .c_in_byte(c_in_byte),
        .ss_o(ss_o), .dbg_ct_idx_i(dbg_ct_idx_i), .dbg_ct_o(dbg_ct_o),
        .dbg_slot_i(dbg_slot_i), .dbg_idx_i(dbg_idx_i), .dbg_coeff_o(dbg_coeff_o),
        .dbg_byte_sel_i(dbg_byte_sel_i), .dbg_byte_idx_i(dbg_byte_idx_i), .dbg_byte_o(dbg_byte_o),
        .dbg_dk_idx_i(dbg_dk_idx_i), .dbg_dk_o(dbg_dk_o),
        .dbg_rho_o(dbg_rho_o), .dbg_sigma_o(dbg_sigma_o),
        .dbg_r_o(dbg_r_o), .dbg_hek_o(dbg_hek_o),
        .dbg_mprime_o(dbg_mprime_o), .dbg_kbar_o(dbg_kbar_o),
        .dbg_decz_o(dbg_decz_o), .dbg_dech_o(dbg_dech_o)
    );
    always #5 clk = ~clk;

    // storage shuttled between operations (the "wire" between Alice and Bob)
    reg [7:0] ek_b  [0:EKB-1];
    reg [7:0] dk_b  [0:DKB-1];
    reg [7:0] ct_b  [0:CTB-1];
    reg [7:0] msg_b [0:MLEN-1];          // "hello world"
    reg [7:0] enc_b [0:MLEN-1];          // XOR-encrypted
    reg [7:0] dec_b [0:MLEN-1];          // XOR-decrypted
    reg [255:0] shared_key, recovered_key;
    integer c, i, j, errors;

    // ---- pulse start and wait for done ----
    task run_op;
        input [1:0] op;
        begin
            op_i = op;
            start_i = 1'b1; @(posedge clk); start_i = 1'b0;
            c = 0;
            while (!done_o && c < 4000000) begin @(posedge clk); c = c + 1; end
            if (!done_o) begin $display("FAIL: op=%0d timeout", op); $finish; end
        end
    endtask

    task print_hex32;        // print a 256-bit value as 32 hex bytes, byte 0 first
        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;
        // message bytes "hello world"
        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 (hardware, ML-KEM-512) ===");
        $write("Original: \"");
        for (i = 0; i < MLEN; i = i + 1) $write("%c", msg_b[i]);
        $display("\"\n");

        rst_n=0; repeat(4) @(posedge clk); rst_n=1; @(posedge clk);

        // ===== Step 1: Alice — KeyGen(d, z) =====
        d_i = {32{8'h42}};     // d = 0x42 repeated (byte 0 in d_i[7:0])
        z_i = {32{8'h77}};     // z = 0x77 repeated
        $display("[1] Alice KeyGen");
        $write("    in  d = "); print_hex32(d_i);
        $write("    in  z = "); print_hex32(z_i);
        run_op(2'd0);
        // read ek (sel=0) out of ek_bram
        dbg_byte_sel_i = 1'b0;
        for (i = 0; i < EKB; i = i + 1) begin
            dbg_byte_idx_i = i[10:0]; @(posedge clk); @(posedge clk);
            ek_b[i] = dbg_byte_o;
        end
        // read full dk (1632 B) via dbg_dk
        for (i = 0; i < DKB; i = i + 1) begin
            dbg_dk_idx_i = i[11:0]; @(posedge clk); @(posedge clk);
            dk_b[i] = dbg_dk_o;
        end
        $display("    out ek = %0d B (ek[0:8] = %02x %02x %02x %02x %02x %02x %02x %02x ...)",
                 EKB, ek_b[0],ek_b[1],ek_b[2],ek_b[3],ek_b[4],ek_b[5],ek_b[6],ek_b[7]);
        $display("    out dk = %0d B (dk[0:8] = %02x %02x %02x %02x %02x %02x %02x %02x ...)",
                 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: Bob — Encaps(ek, m) =====
        // stream ek back into ek_bram (Bob receives Alice's public key)
        for (i = 0; i < EKB; i = i + 1) begin
            ek_in_we = 1'b1; ek_in_addr = i[10:0]; ek_in_byte = ek_b[i];
            @(posedge clk);
        end
        ek_in_we = 1'b0; @(posedge clk);
        msg_i = {32{8'hDE}};   // m = 0xDE repeated
        $display("[2] Bob Encaps");
        $write("    in  ek = %0d B;  in m = ", EKB); print_hex32(msg_i);
        run_op(2'd1);
        shared_key = ss_o;
        // read ct (768 B) out of ct_bram
        for (i = 0; i < CTB; i = i + 1) begin
            dbg_ct_idx_i = i[10:0]; @(posedge clk); @(posedge clk);
            ct_b[i] = dbg_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 "hello world" with shared_key =====
        for (i = 0; i < MLEN; i = i + 1)
            enc_b[i] = msg_b[i] ^ shared_key[8*(i % 32) +: 8];
        $display("[3] Bob encrypt (XOR stream, demo)");
        $write("    in  msg       = "); for (i=0;i<MLEN;i=i+1) $write("%02x ", msg_b[i]); $write("\n");
        $write("    out encrypted = "); for (i=0;i<MLEN;i=i+1) $write("%02x ", enc_b[i]); $write("\n\n");

        // ===== Step 4: Alice — Decaps(dk, kem_ct) =====
        for (i = 0; i < DKB; i = i + 1) begin
            dk_in_we = 1'b1; dk_in_addr = i[11:0]; dk_in_byte = dk_b[i];
            @(posedge clk);
        end
        dk_in_we = 1'b0;
        for (i = 0; i < CTB; i = i + 1) begin
            c_in_we = 1'b1; c_in_addr = i[10:0]; c_in_byte = ct_b[i];
            @(posedge clk);
        end
        c_in_we = 1'b0; @(posedge clk);
        $display("[4] Alice Decaps");
        $display("    in  dk = %0d B;  in kem_ct = %0d B", DKB, CTB);
        run_op(2'd2);
        recovered_key = ss_o;
        $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];
        $display("[5] Alice decrypt");
        $write("    in  encrypted = "); for (i=0;i<MLEN;i=i+1) $write("%02x ", enc_b[i]); $write("\n");
        $write("    out decrypted = \"");
        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 != recovered_key");
            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: keys match, message recovered. (hardware ML-KEM-512)");
        else
            $display("FAILURE: %0d mismatches", errors);
        $finish;
    end

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