`timescale 1ns / 1ps

module cam_core #(
    parameter int NUM_ROWS  = 512,
    parameter int HASH_BITS = 512,
    parameter int LANES     = 16,

    parameter int ROW_BITS   = $clog2(NUM_ROWS),
    parameter int SCORE_BITS = $clog2(HASH_BITS + 1)
) (
    input logic clk,
    input logic rst_n,

    // Static load interface. In the first prototype, writes are expected
    // before online queries begin.
    input logic [ ROW_BITS-1:0] wr_row,
    input logic [HASH_BITS-1:0] wr_hash,
    input logic                 wr_en,

    // Single-request blocking query interface.
    input  logic                 query_valid,
    output logic                 query_ready,
    input  logic [HASH_BITS-1:0] query_hash,

    output logic                  result_valid,
    input  logic                  result_ready,
    output logic [  ROW_BITS-1:0] top1_index,
    output logic [SCORE_BITS-1:0] top1_score,

`ifdef SIM_NOISE
    // Flattened for easier cocotb/Verilator handling:
    // lane k mask = noise_mask_lanes_flat[k*HASH_BITS +: HASH_BITS]
    input logic [LANES*HASH_BITS-1:0] noise_mask_lanes_flat,
`endif

`ifdef SIM_DEBUG
    // Flattened for easier cocotb/Verilator handling:
    // score[row] = score_debug_flat[row*SCORE_BITS +: SCORE_BITS]
    output logic [NUM_ROWS*SCORE_BITS-1:0] score_debug_flat,
`endif

    output logic busy
);
    localparam int NUM_BATCHES = (NUM_ROWS + LANES - 1) / LANES;

    typedef enum logic [1:0] {
        S_IDLE,
        S_SCAN,
        S_DONE
    } state_t;

    state_t state_q, state_d;

    logic [HASH_BITS-1:0] cam_mem[NUM_ROWS];

    logic [HASH_BITS-1: 0]    query_q;
    logic [ROW_BITS-1 : 0]    base_row_q;
    logic [ROW_BITS-1 : 0]    base_row_d;

    logic [ROW_BITS-1:0] best_index_q, best_index_d;
    logic [SCORE_BITS-1:0] best_score_q, best_score_d;

    logic [  ROW_BITS-1:0] lane_best_index;
    logic [SCORE_BITS-1:0] lane_best_score;

    logic [  ROW_BITS-1:0] lane_best_index_next[LANES+1];
    logic [SCORE_BITS-1:0] lane_best_score_next[LANES+1];

    logic [SCORE_BITS-1:0] lane_score          [  LANES];
    logic [  ROW_BITS-1:0] lane_row            [  LANES];
    logic                  lane_valid          [  LANES];

    assign query_ready  = (state_q == S_IDLE);
    assign result_valid = (state_q == S_DONE);
    assign top1_index   = best_index_q;
    assign top1_score   = best_score_q;
    assign busy         = (state_q == S_SCAN);

    // Memory write path.
    always_ff @(posedge clk) begin
        if (wr_en) begin
            cam_mem[wr_row] <= wr_hash;
        end
    end

`ifdef SIM_DEBUG
    // Clear by default. Individual rows are overwritten during scan.
    initial begin
        score_debug_flat = '0;
    end
`endif

    genvar lane;
    generate
        for (lane = 0; lane < LANES; lane++) begin : gen_lanes
            logic [HASH_BITS-1:0] row_hash;
            logic [HASH_BITS-1:0] effective_hash;
            logic [HASH_BITS-1:0] match_bits;

            assign lane_row[lane] = base_row_q + lane[ROW_BITS-1:0];
            assign lane_valid[lane] = (lane_row[lane] < NUM_ROWS[ROW_BITS-1:0]);

            // This read is modeled behaviorally. Later this can be replaced
            // by banked BRAM/URAM without changing the compare path.
            assign row_hash = lane_valid[lane] ? cam_mem[lane_row[lane]] : '0;

`ifdef SIM_NOISE
            logic [HASH_BITS-1:0] lane_noise_mask;
            assign lane_noise_mask = noise_mask_lanes_flat[lane*HASH_BITS+:HASH_BITS];
            assign effective_hash  = row_hash ^ lane_noise_mask;
`else
            assign effective_hash = row_hash;
`endif

            assign match_bits = ~(query_q ^ effective_hash);

            popcount #(
                .WIDTH(HASH_BITS),
                .GROUP(8),
                .OUT_WIDTH(SCORE_BITS)
            ) u_popcount (
                .bits_i (match_bits),
                .count_o(lane_score[lane])
            );

            argmax_update #(
                .ROW_BITS  (ROW_BITS),
                .SCORE_BITS(SCORE_BITS)
            ) u_argmax_lane (
                .best_index_i(lane_best_index_next[lane]),
                .best_score_i(lane_best_score_next[lane]),
                .cand_index_i(lane_row[lane]),
                .cand_score_i(lane_valid[lane] ? lane_score[lane] : '0),
                .best_index_o(lane_best_index_next[lane+1]),
                .best_score_o(lane_best_score_next[lane+1])
            );
        end
    endgenerate

    assign lane_best_index_next[0] = best_index_q;
    assign lane_best_score_next[0] = best_score_q;
    assign lane_best_index = lane_best_index_next[LANES];
    assign lane_best_score = lane_best_score_next[LANES];

    always_comb begin
        state_d      = state_q;
        base_row_d   = base_row_q;
        best_index_d = best_index_q;
        best_score_d = best_score_q;

        unique case (state_q)
            S_IDLE: begin
                if (query_valid) begin
                    state_d      = S_SCAN;
                    base_row_d   = '0;
                    best_index_d = {ROW_BITS{1'b1}};
                    best_score_d = '0;
                end
            end

            S_SCAN: begin
                best_index_d = lane_best_index;
                best_score_d = lane_best_score;

                if (base_row_q + LANES >= NUM_ROWS) begin
                    state_d = S_DONE;
                end else begin
                    base_row_d = base_row_q + LANES[ROW_BITS-1:0];
                end
            end

            S_DONE: begin
                if (result_ready) begin
                    state_d = S_IDLE;
                end
            end

            default: begin
                state_d = S_IDLE;
            end
        endcase
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_q      <= S_IDLE;
            query_q      <= '0;
            base_row_q   <= '0;
            best_index_q <= '0;
            best_score_q <= '0;
        end else begin
            state_q      <= state_d;
            base_row_q   <= base_row_d;
            best_index_q <= best_index_d;
            best_score_q <= best_score_d;

            if ((state_q == S_IDLE) && query_valid) begin
                query_q <= query_hash;
            end

`ifdef SIM_DEBUG
            if (state_q == S_IDLE && query_valid) begin
                score_debug_flat <= '0;
            end else if (state_q == S_SCAN) begin
                for (int l = 0; l < LANES; l++) begin
                    if (lane_valid[l]) begin
                        score_debug_flat[lane_row[l]*SCORE_BITS+:SCORE_BITS] <= lane_score[l];
                    end
                end
            end
`endif
        end
    end
endmodule