// tb_kg_xsim.v - KeyGen-only KAT testbench for mlkem_top
//
// Runs full ML-KEM flow (kg→en→de) but only CHECKS KeyGen results.
// Uses existing vector files from sync_rtl/top/TB/vectors/.
//
// Usage:
//   xvlog -sv -i . <all_deps>.v sync_rtl/kg/TB/tb_kg_xsim.v
//   xelab tb_kg_xsim -s tb_kg_xsim --timescale 1ns/1ps
//   xsim tb_kg_xsim -R

`timescale 1ns / 1ps

module tb_kg_xsim;

    parameter VECTOR_FILE   = "sync_rtl/top/TB/vectors/mlkem_top_input.hex";
    parameter EXPECTED_FILE = "sync_rtl/top/TB/vectors/mlkem_top_expected.hex";
    parameter MAX_VECTORS   = 16;
    parameter TIMEOUT_CYCLES = 10000000;
    parameter K_PARAM       = 4;
    localparam PK_WIDTH      = 12 * K_PARAM * 256;
    localparam SK_WIDTH      = 12 * K_PARAM * 256;
    localparam EXP_PK_WIDTH  = 6400;
    localparam EXP_SK_WIDTH  = 13056;
    localparam EXP_CT_WIDTH  = 6144;
    localparam EXP_SS_WIDTH  = 256;

    // DUT signals
    reg             clk;
    reg             rst_n;
    reg  [1:0]      mode;
    reg  [2:0]      i_k;
    reg             valid_i;
    wire            ready_o;
    wire [PK_WIDTH-1:0]  pk_o;
    wire [SK_WIDTH-1:0]  sk_o;
    wire                 pk_valid;
    wire                 sk_valid;
    wire [EXP_CT_WIDTH*K_PARAM/2-1:0] ct_o;
    wire [255:0]         K_o;
    wire                 ct_valid;
    wire                 K_valid;
    wire [255:0]         K_o_dec;
    wire                 K_valid_dec;
    wire                 done_o;

    // DUT instantiation
    mlkem_top #(.K(K_PARAM)) u_dut (
        .clk(clk), .rst_n(rst_n), .mode(mode), .i_k(i_k),
        .valid_i(valid_i), .ready_o(ready_o),
        .pk_o(pk_o), .sk_o(sk_o), .pk_valid(pk_valid), .sk_valid(sk_valid),
        .ct_o(ct_o), .K_o(K_o), .ct_valid(ct_valid), .K_valid(K_valid),
        .K_o_dec(K_o_dec), .K_valid_dec(K_valid_dec), .done_o(done_o)
    );

    initial clk = 1'b0;
    always #5 clk = ~clk;

    reg [767:0]   input_mem   [0:MAX_VECTORS-1];
    reg [25855:0] expected_mem [0:MAX_VECTORS-1];

    integer vec_count, idx, kg_pass, kg_fail;

    function [7:0] nibble_to_ascii;
        input [3:0] nibble;
        nibble_to_ascii = (nibble < 4'd10) ? (8'h30 + {4'd0, nibble}) : (8'h41 + ({4'd0, nibble} - 4'd10));
    endfunction

    task print_hex256;
        input [255:0] val;
        input [256*8:1] label;
        integer bi;
        begin
            $write("%s: ", label);
            for (bi = 63; bi >= 0; bi = bi - 1)
                $write("%c", nibble_to_ascii(val[(bi*4)+:4]));
            $write("\n");
        end
    endtask

    integer wfd_result;
    task wait_for_done;
        input [256*8:1] op_name;
        integer cyc;
        begin
            cyc = 0;
            while (!done_o && cyc < TIMEOUT_CYCLES) begin @(posedge clk); cyc = cyc + 1; end
            if (cyc >= TIMEOUT_CYCLES) begin
                $display("ERROR: %s timeout after %0d cycles", op_name, TIMEOUT_CYCLES);
                wfd_result = 0;
            end else begin
                $display("INFO:  %s done after %0d cycles", op_name, cyc);
                wfd_result = 1;
            end
        end
    endtask

    initial begin
        vec_count = 0;
        $readmemh(VECTOR_FILE, input_mem);
        $readmemh(EXPECTED_FILE, expected_mem);

        begin
            integer found_end = 0;
            for (idx = 0; idx < MAX_VECTORS; idx = idx + 1) begin
                if (!found_end && (input_mem[idx] === 768'hx || input_mem[idx] === 768'hz))
                    found_end = 1;
                else if (!found_end)
                    vec_count = vec_count + 1;
            end
        end

        if (vec_count == 0) begin
            $display("ERROR: No vectors loaded from %s", VECTOR_FILE);
            $finish;
        end

        $display("====================================================");
        $display("MLKEM_TOP KeyGen-ONLY KAT TESTBENCH");
        $display("  Vectors loaded: %0d", vec_count);
        $display("====================================================");

        mode    <= 2'd0;
        i_k     <= 3'd2;
        valid_i <= 1'b0;
        rst_n   <= 1'b0;
        repeat (5) @(posedge clk);
        rst_n   <= 1'b1;
        @(posedge clk);

        force u_dut.chain_kc_ready_o = 1'b1;
        force u_dut.ntt_valid_o = 1'b1;

        kg_pass = 0; kg_fail = 0;

        for (idx = 0; idx < vec_count; idx = idx + 1) begin
            reg [255:0] d_val, msg_val, z_val;
            reg [EXP_PK_WIDTH-1:0]  exp_pk;
            reg [EXP_SK_WIDTH-1:0]  exp_sk;

            d_val   = input_mem[idx][767:512];
            msg_val = input_mem[idx][511:256];
            z_val   = input_mem[idx][255:0];
            exp_sk = expected_mem[idx][EXP_SS_WIDTH + EXP_CT_WIDTH +: EXP_SK_WIDTH];
            exp_pk = expected_mem[idx][EXP_SS_WIDTH + EXP_CT_WIDTH + EXP_SK_WIDTH +: EXP_PK_WIDTH];

            $display("--- Vector %0d: KeyGen ---", idx);
            print_hex256(d_val, "  d  ");

            // KeyGen
            force u_dut.d_reg = d_val;
            mode    <= 2'b00;
            i_k     <= 3'd2;
            valid_i <= 1'b1;
            @(posedge clk);
            valid_i <= 1'b0;
            wait_for_done("KeyGen");
            if (wfd_result) begin
                release u_dut.d_reg;
                if (pk_valid && sk_valid) begin
                    if (pk_o[EXP_PK_WIDTH-1:0] == exp_pk) begin
                        $display("  PASS: pk matches expected");
                        kg_pass = kg_pass + 1;
                    end else begin
                        $display("  FAIL: pk mismatch");
                        kg_fail = kg_fail + 1;
                    end
                    if (sk_o[SK_WIDTH-1:0] == exp_sk[SK_WIDTH-1:0])
                        $display("  PASS: sk matches expected");
                    else begin
                        $display("  FAIL: sk mismatch");
                        kg_fail = kg_fail + 1;
                    end
                end else begin
                    $display("  FAIL: pk_valid=%b sk_valid=%b", pk_valid, sk_valid);
                    kg_fail = kg_fail + 1;
                end
            end else begin
                release u_dut.d_reg;
                $display("  FAIL: KeyGen timeout");
                kg_fail = kg_fail + 1;
                rst_n <= 1'b0; repeat (5) @(posedge clk); rst_n <= 1'b1;
                force u_dut.chain_kc_ready_o = 1'b1; force u_dut.ntt_valid_o = 1'b1;
                @(posedge clk);
            end
            repeat (2) @(posedge clk);

            // Encaps (run but don't check)
            force u_dut.m_reg = msg_val;
            mode    <= 2'b01;
            i_k     <= 3'd2;
            valid_i <= 1'b1;
            @(posedge clk);
            valid_i <= 1'b0;
            wait_for_done("Encaps");
            if (wfd_result) release u_dut.m_reg;
            else begin
                release u_dut.m_reg;
                rst_n <= 1'b0; repeat (5) @(posedge clk); rst_n <= 1'b1;
                force u_dut.chain_kc_ready_o = 1'b1; force u_dut.ntt_valid_o = 1'b1;
                @(posedge clk);
            end
            repeat (2) @(posedge clk);

            // Decaps (run but don't check)
            force u_dut.z_reg = z_val;
            mode    <= 2'b10;
            i_k     <= 3'd2;
            valid_i <= 1'b1;
            @(posedge clk);
            valid_i <= 1'b0;
            wait_for_done("Decaps");
            if (wfd_result) release u_dut.z_reg;
            else begin
                release u_dut.z_reg;
                rst_n <= 1'b0; repeat (5) @(posedge clk); rst_n <= 1'b1;
                force u_dut.chain_kc_ready_o = 1'b1; force u_dut.ntt_valid_o = 1'b1;
                @(posedge clk);
            end
            repeat (2) @(posedge clk);
        end

        release u_dut.chain_kc_ready_o;
        release u_dut.ntt_valid_o;

        $display("====================================================");
        $display("KEYGEN TEST COMPLETE");
        $display("  PASS: %0d  FAIL: %0d", kg_pass, kg_fail);
        $display("====================================================");
        $finish;
    end

    initial begin
        #(TIMEOUT_CYCLES * 10 * 100);
        $display("FATAL: Global simulation timeout");
        $finish;
    end

endmodule