`timescale 1ns / 1ps
`include "cam_params.svh"

module match_engine_pipeline (
    input  logic                              clk,
    input  logic                              rst_n,
    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]            result_row,
    output logic [(`SCORE_BITS)-1:0]          result_score,
    output logic                              busy,
    output logic                              rd_valid_o,
    output logic [(`ROW_BITS)-1:0]            rd_base_row_o,
    input  logic                              rd_valid_i,
    input  logic [(`LANES)*(`ROW_BITS)-1:0]   rd_row_ids_i,
    input  logic [(`LANES)*(`HASH_BITS)-1:0]  rd_hashes_i,
    input  logic [(`LANES)-1:0]               rd_lane_valid_i
`ifdef SIM_DEBUG
    ,output logic [(`NUM_ROWS)*(`SCORE_BITS)-1:0] score_debug_flat
`endif
);
    typedef enum logic [1:0] {S_IDLE, S_SCAN, S_DRAIN, S_DONE} state_t;
    state_t state_q;

    logic [(`HASH_BITS)-1:0] query_q;
    logic [(`ROW_BITS)-1:0] issue_base_q;
    logic [$clog2(`NUM_ROWS/`LANES+1)-1:0] returned_q;
    logic [(`ROW_BITS)-1:0] best_row_q;
    logic [(`SCORE_BITS)-1:0] best_score_q;

    logic [(`HASH_BITS)-1:0] match_bits [0:`LANES-1];
    logic score_valid [0:`LANES-1];
    logic [(`ROW_BITS)-1:0] score_row [0:`LANES-1];
    logic [(`SCORE_BITS)-1:0] lane_score [0:`LANES-1];
    logic [(`ROW_BITS)-1:0] batch_row;
    logic [(`SCORE_BITS)-1:0] batch_score;

    assign query_ready  = (state_q == S_IDLE);
    assign result_valid = (state_q == S_DONE);
    assign result_row   = best_row_q;
    assign result_score = best_score_q;
    assign busy         = (state_q != S_IDLE);
    assign rd_valid_o   = (state_q == S_SCAN);
    assign rd_base_row_o = issue_base_q;

    generate
        for (genvar lane = 0; lane < `LANES; lane++) begin : gen_scores
            assign match_bits[lane] = ~(query_q ^ rd_hashes_i[lane*`HASH_BITS +: `HASH_BITS]);
            popcount_pipeline #(
                .WIDTH(`HASH_BITS),
                .ROW_BITS(`ROW_BITS),
                .OUT_WIDTH(`SCORE_BITS)
            ) u_popcount_pipeline (
                .clk(clk),
                .rst_n(rst_n),
                .valid_i(rd_valid_i && rd_lane_valid_i[lane]),
                .row_i(rd_row_ids_i[lane*`ROW_BITS +: `ROW_BITS]),
                .bits_i(match_bits[lane]),
                .valid_o(score_valid[lane]),
                .row_o(score_row[lane]),
                .count_o(lane_score[lane])
            );
        end
    endgenerate

    always_comb begin
        batch_row = score_row[0];
        batch_score = lane_score[0];
        for (int lane = 1; lane < `LANES; lane++) begin
            if ((lane_score[lane] > batch_score) ||
                ((lane_score[lane] == batch_score) && (score_row[lane] < batch_row))) begin
                batch_score = lane_score[lane];
                batch_row = score_row[lane];
            end
        end
    end

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_q      <= S_IDLE;
            query_q      <= '0;
            issue_base_q <= '0;
            returned_q   <= 0;
            best_row_q   <= '0;
            best_score_q <= '0;
`ifdef SIM_DEBUG
            score_debug_flat <= '0;
`endif
        end else begin
            unique case (state_q)
                S_IDLE: begin
                    if (query_valid) begin
                        query_q      <= query_hash;
                        issue_base_q <= '0;
                        returned_q   <= 0;
                        best_row_q   <= TIE_BREAK_SENTINEL;
                        best_score_q <= '0;
`ifdef SIM_DEBUG
                        score_debug_flat <= '0;
`endif
                        state_q <= S_SCAN;
                    end
                end
                S_SCAN: begin
                    if (issue_base_q + `LANES >= `NUM_ROWS) begin
                        state_q <= S_DRAIN;
                    end else begin
                        issue_base_q <= issue_base_q + `LANES;
                    end
                end
                S_DRAIN: begin
                    if ((returned_q == (`NUM_ROWS / `LANES)) && !score_valid[0]) begin
                        state_q <= S_DONE;
                    end
                end
                S_DONE: begin
                    if (result_ready) state_q <= S_IDLE;
                end
                default: state_q <= S_IDLE;
            endcase

            if (score_valid[0]) begin
                returned_q <= returned_q + 1'b1;
                if ((batch_score > best_score_q) ||
                    ((batch_score == best_score_q) && (batch_row < best_row_q))) begin
                    best_score_q <= batch_score;
                    best_row_q <= batch_row;
                end
`ifdef SIM_DEBUG
                for (int lane = 0; lane < `LANES; lane++) begin
                    score_debug_flat[score_row[lane]*`SCORE_BITS +: `SCORE_BITS] <= lane_score[lane];
                end
`endif
            end
        end
    end
endmodule