`timescale 1ns/1ps

module TB;

    // ====================================================
    // Parameters
    // ====================================================
    parameter FIFO_DEPTH       = 64;
    parameter SCRAMBLER_SEED   = 32'hFFFFFFFF;
    parameter CLK_PERIOD       = 10;          // 10ns
    parameter PATN_COUNT       = 20'd10;      // Training match count
    parameter TAP_STEP         = 3'd3;        // Delay adjustment step

    // ====================================================
    // Constants (moved to top for visibility)
    // ====================================================
    localparam [31:0] TRAINING_PATN = 32'h68666E6C;  // "hfnl"
    localparam [31:0] TRAINING_EXIT = 32'h65786974;  // "exit"
    localparam [31:0] FRAME_HEADER  = 32'hBCBCBCBC;

initial begin
  $fsdbAutoSwitchDumpfile(500, "./verdplus.fsdb", 1000000);
  $fsdbDumpvars();
  $fsdbDumpMDA();
end 

    // ====================================================
    // Signals
    // ====================================================
    logic clk;
    logic rst_n;
    logic [3:0] serial_in;
    logic [19:0] patn_count;
    logic [2:0] tap_step;
    logic link_down;
    logic [2:0] delay_tap;
    logic [12:0] wr_addr;
    logic [511:0] wr_data;
    logic wr_en;
    logic [63:0] byte_mask;
    logic crc_error;
    logic [2  :0] tap_adj_mask = 3'b111;      // Delay adjustment mask
    logic [0  :0] tap_force = 1'b0   ;        // Delay force to tap_step
    logic         tap_adj_req  ;
    logic         frame_done   ;
    logic [31 :0] train_status ;
    logic [31 :0] frame_status ;
    logic         always_on = 1'b0  ; 
    logic         prefilling   ;
    logic         prefill_start;
    
    // Added for scrambler test
    logic         descram_en;
    logic [31:0]  lfsr_lane[0:3];   // 4 independent LFSRs for scrambling
    logic         train_ready;
    logic         force_train;        // Force Training
    // ====================================================
    // DUT Instance
    // ====================================================
    ulink_rx #(
      .FIFO_DEPTH     ( FIFO_DEPTH     )
     ,.SCRAMBLER_SEED ( SCRAMBLER_SEED )
    ) dut (
       .clk           ( clk            )
      ,.rst_n         ( rst_n          )
      ,.serial_in     ( serial_in      )
      ,.patn_count    ( patn_count     )
      ,.tap_step      ( tap_step       )
      ,.descram_en    ( descram_en     )   // now variable
      ,.link_down     ( link_down      )
      ,.delay_tap     ( delay_tap      )
      ,.wr_addr       ( wr_addr        )
      ,.wr_data       ( wr_data        )
      ,.wr_en         ( wr_en          )
      ,.byte_mask     ( byte_mask      )
      ,.crc_error     ( crc_error      )
      ,.tap_adj_mask  ( tap_adj_mask   )
      ,.tap_force     ( tap_force      )
      ,.tap_adj_req   ( tap_adj_req    )
      ,.frame_done    ( frame_done     )
      ,.train_status  ( train_status   )
      ,.frame_status  ( frame_status   )
      ,.always_on     ( always_on      )
      ,.prefilling    ( prefilling     )
      ,.prefill_start ( prefill_start  )
	  ,.train_ready   ( train_ready    )
	  ,.force_train   ( force_train    )
    );

    // ====================================================
    // Clock Generation
    // ====================================================
    initial begin
        clk = 0;
        forever #(CLK_PERIOD/2) clk = ~clk;
    end

    // ====================================================
    // Tasks for Stimulus Generation
    // ====================================================

    // Reset task
    task reset;
        begin
            rst_n = 0;
            repeat(5) @(posedge clk);
            rst_n = 1;
            repeat(2) @(posedge clk);
        end
    endtask

    // Send one 128-bit block over 32 cycles
    task send_block;
        input [127:0] block;
        begin
            integer i;
            for (i = 31; i >= 0; i = i - 1) begin
                serial_in[0] = block[i];
                serial_in[1] = block[32 + i];
                serial_in[2] = block[64 + i];
                serial_in[3] = block[96 + i];
                @(posedge clk);
                #0.1;
            end
        end
    endtask

    // Send multiple identical blocks
    task send_blocks;
        input [127:0] block;
        input integer count;
        begin
            repeat(count) send_block(block);
        end
    endtask

    // Send a sequence of 32-bit words (packed into 128-bit blocks)
    task send_words;
        input integer num_words;
        input [31:0] words[0:255];
        begin
            integer word_idx;
            integer j;
            reg [127:0] block;
            word_idx = 0;
            while (word_idx < num_words) begin
                block = 128'h0;
                for (j = 0; j < 4; j = j + 1) begin
                    if (word_idx < num_words) begin
                        block[32*j +: 32] = words[word_idx];
                        word_idx = word_idx + 1;
                    end
                end
                send_block(block);
            end
        end
    endtask

   // ====================================================
// Scrambler tasks (used when descram_en=1)
// ====================================================

// Original scrambler (all lanes scrambled)
task scramble_block;
    input [127:0] original;
    output [127:0] scrambled;
    integer i;
    reg [31:0] lane_word;
    begin
        scrambled = 128'h0;
        for (i = 0; i < 4; i = i + 1) begin
            lane_word = original[32*i +: 32];
            scrambled[32*i +: 32] = lane_word ^ lfsr_lane[i];
            // Update LFSR: feedback = lfsr[31]^lfsr[30]^lfsr[29]^lfsr[28]
            lfsr_lane[i] = {lfsr_lane[i][30:0], 
                             lfsr_lane[i][31] ^ lfsr_lane[i][30] ^ 
                             lfsr_lane[i][29] ^ lfsr_lane[i][28]};
        end
    end
endtask

// New scrambler with per-lane mask (1=scramble that lane)
task scramble_block_masked;
    input [127:0] original;
    input [3:0]   mask;        // 1=scramble this lane
    output [127:0] scrambled;
    integer i;
    reg [31:0] lane_word;
    begin
        scrambled = 128'h0;
        for (i = 0; i < 4; i = i + 1) begin
            lane_word = original[32*i +: 32];
            if (mask[i]) begin
                scrambled[32*i +: 32] = lane_word ^ lfsr_lane[i];
                // Update LFSR only for scrambled lanes
                lfsr_lane[i] = {lfsr_lane[i][30:0], 
                                 lfsr_lane[i][31] ^ lfsr_lane[i][30] ^ 
                                 lfsr_lane[i][29] ^ lfsr_lane[i][28]};
            end else begin
                scrambled[32*i +: 32] = lane_word;
                // No LFSR update for unscrambled lanes
            end
        end
    end
endtask

    // ====================================================
    // CRC32 Computation Functions
    // ====================================================
    function automatic [31:0] crc32_next;
        input [31:0] crc;
        input [31:0] data;
        reg [31:0] new_crc;
        reg b;
        integer i;
        begin
            new_crc = crc;
            for (i = 0; i < 32; i = i + 1) begin
                b = new_crc[31] ^ data[31-i];
                new_crc = {new_crc[30:0], 1'b0};
                if (b) new_crc = new_crc ^ 32'h04C11DB7;
            end
            crc32_next = new_crc;
        end
    endfunction

    function automatic [31:0] crc32_compute;
        input [31:0] words[0:255];
        input integer num;
        reg [31:0] crc;
        integer i;
        begin
            crc = 32'hFFFFFFFF;
            for (i = 0; i < num; i = i + 1) begin
                crc = crc32_next(crc, words[i]);
            end
            crc32_compute = crc;
        end
    endfunction

    // ====================================================
    // Test Variables (module level)
    // ====================================================
    integer data_words;
    reg [31:0] frame_words[0:255];
    integer frame_len;
    integer i;
    // ====================================================
    // Test Sequence
    // ====================================================
    initial begin
        // Initialize
        serial_in = 4'h0;
        patn_count = PATN_COUNT;
        tap_step = TAP_STEP;
        descram_en = 0;   // default disabled
        force_train = 0;

        $display("========================================");
        $display("Testbench started at %0t", $time);
        $display("========================================");

        reset();
        wait(prefill_start == 1);
        repeat(10) @(posedge clk);
        prefilling = 1;
        repeat(100) @(posedge clk);
        prefilling = 0;
        // -------------------------------------------------
        // Phase 1: Link Training (successful)
        // -------------------------------------------------
        $display("Phase 1: Training with correct patterns...");

        send_blocks({4{TRAINING_PATN}}, PATN_COUNT - 0);
        send_blocks({4{TRAINING_EXIT}}, 1);

        wait(dut.u_train.train_ready == 1);
        $display("Link ready at %0t", $time);

        // -------------------------------------------------
        // Phase 2: Send a correct frame
        // -------------------------------------------------
        $display("Phase 2: Sending a correct frame...");

        data_words = 20;
        frame_len = 2 + data_words + 1;

        frame_words[0] = FRAME_HEADER;
        frame_words[1] = (32'h1235 << 16) | data_words;

        for (i = 0; i < data_words; i = i + 1) begin
            frame_words[2+i] = 32'hA0000000 + i;
        end

        // Compute CRC over addr_len and data (1 + data_words words)
        begin
            reg [31:0] crc_tmp[0:255];
            integer k;
            for (k = 0; k < 1 + data_words; k = k + 1) begin
                crc_tmp[k] = frame_words[1 + k];
            end
            frame_words[2 + data_words] = crc32_compute(crc_tmp, 1 + data_words);
        end

        send_words(frame_len, frame_words);

        fork
            begin
                @(posedge wr_en);
                $display("Write detected: addr=%0d data=%h mask=%h", wr_addr, wr_data, byte_mask);
                if (wr_addr !== 13'h123) $error("Unexpected write address: %0d", wr_addr);
                if (byte_mask[31:0] !== 32'hFFFF_0000) $error("Byte mask mismatch: %h", byte_mask);
                if (byte_mask[63:32] !== 32'hFFFF_FFFF) $error("Byte mask high part not zero: %h", byte_mask[63:32]);
                $display("Correct frame write verified.");
            end
        join_none

        repeat(100) @(posedge clk);

        // -------------------------------------------------
        // Phase 3: Send a frame with bad CRC
        // -------------------------------------------------
        $display("Phase 3: Sending a frame with bad CRC...");

        frame_words[2+data_words] = ~frame_words[2+data_words];
        send_words(frame_len, frame_words);

        @(posedge crc_error);
        $display("CRC error detected at %0t", $time);
        repeat(10) @(posedge clk);
        if (link_down) $display("Link down as expected.");
        else $error("Link not down after CRC error.");

        // -------------------------------------------------
        // Phase 4: Re-train and send another correct frame
        // -------------------------------------------------
        $display("Phase 4: Re-training...");

        send_blocks({4{TRAINING_PATN}}, PATN_COUNT - 0);
        send_blocks({4{TRAINING_EXIT}}, 1);

        wait(dut.u_train.train_ready == 1);
        $display("Link ready again.");

        // Compute CRC over addr_len and data (1 + data_words words)
        begin
            reg [31:0] crc_tmp[0:255];
            integer k;
            for (k = 0; k < 1 + data_words; k = k + 1) begin
                crc_tmp[k] = frame_words[1 + k];
            end
            frame_words[2 + data_words] = crc32_compute(crc_tmp, 1 + data_words);
        end
        send_words(frame_len, frame_words);

        repeat(100) @(posedge clk);

        // -------------------------------------------------
        // Phase 5: Test delay_tap adjustment on match failure
        // -------------------------------------------------
        $display("Phase 5: Testing delay_tap adjustment...");

        force_train = 1;
		repeat(1) @(posedge clk);
		force_train = 0;
        send_blocks(128'h0, 5);
        repeat(200) @(posedge clk);
        $display("Final delay_tap = %0d", delay_tap);

        // -------------------------------------------------
        // Phase 6: Test with scrambler enabled (header not scrambled)
        // -------------------------------------------------
        $display("Phase 6: Testing with descrambler enabled (header not scrambled)...");

        // Reset and train with descram_en=0 (training patterns are not scrambled)
        descram_en = 0;
        force_train = 1;
		repeat(1) @(posedge clk);
		force_train = 0;
        send_blocks({4{TRAINING_PATN}}, PATN_COUNT - 0);
        send_blocks({4{TRAINING_EXIT}}, 1);

        wait(dut.u_train.train_ready == 1);
        $display("Link ready for scrambled data.");

        // Enable descrambler in DUT and initialize LFSRs for transmission
        descram_en = 1;
        for (i = 0; i < 4; i = i + 1) lfsr_lane[i] = SCRAMBLER_SEED;

        // Construct a simple frame (small data length)
        data_words = 10;
        frame_len = 2 + data_words + 1;

        frame_words[0] = FRAME_HEADER;
        frame_words[1] = (32'h1234 << 16) | data_words;  // base_addr=0x1234, offset=0, len=10

        for (i = 0; i < data_words; i = i + 1) begin
            frame_words[2+i] = 32'hB0000000 + i;  // different pattern from Phase 2
        end

        // Compute CRC over addr_len and data (1 + data_words words)
        begin
            reg [31:0] crc_tmp[0:255];
            integer k;
            for (k = 0; k < 1 + data_words; k = k + 1) begin
                crc_tmp[k] = frame_words[1 + k];
            end
            frame_words[2 + data_words] = crc32_compute(crc_tmp, 1 + data_words);
        end

        // Now scramble each 128-bit block and send
        begin
            integer word_idx;
            integer j;
            integer block_num;
            reg [127:0] orig_block, scrambled_block;
            reg [3:0]   block_mask;
            word_idx = 0;
            block_num = 0;
            while (word_idx < frame_len) begin
                orig_block = 128'h0;
                // Fill block with up to 4 words, pad with 0 if needed
                for (j = 0; j < 4; j = j + 1) begin
                    if (word_idx < frame_len) begin
                        orig_block[32*j +: 32] = frame_words[word_idx];
                        word_idx = word_idx + 1;
                    end else begin
                        orig_block[32*j +: 32] = 32'h0;  // pad with 0
                    end
                end
                // Determine mask: first block leaves lane0 unscrambled (header), others all scrambled
                block_mask = (block_num == 0) ? 4'b1110 : 4'b1111;
                scramble_block_masked(orig_block, block_mask, scrambled_block);
                send_block(scrambled_block);
                block_num = block_num + 1;
            end
        end

        // Monitor writes
        fork
            begin
                @(posedge wr_en);
                $display("Scrambled test write: addr=%0d data=%h mask=%h", wr_addr, wr_data, byte_mask);
            end
        join_none

        repeat(200) @(posedge clk);

        // Check that no CRC error happened
        if (crc_error) $error("CRC error in scrambled test");
        else $display("Scrambled test passed (no CRC error).");

        // Disable scrambler for rest of test
        descram_en = 0;

        // -------------------------------------------------
        // End of test
        // -------------------------------------------------
        repeat(50) @(posedge clk);
        $display("========================================");
        $display("Testbench finished at %0t", $time);
        $display("========================================");
        $finish;
    end

    // ====================================================
    // Monitor and Assertions
    // ====================================================
    always @(posedge clk) begin
        if (wr_en) $display("WRITE: addr=%0d data=%h mask=%h", wr_addr, wr_data, byte_mask);
        if (crc_error) $display("CRC_ERROR pulse at %0t", $time);
    end

    always @(posedge clk) begin
        if (dut.u_train.train_ready && link_down)
            $error("Inconsistent state: train_ready and link_down both high");
    end

endmodule