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

module match_engine (
    input  logic                                clk,
    input  logic                                rst_n,
    // 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]              result_row,
    output logic [(`SCORE_BITS)-1:0]            result_score,
    output logic                                busy,
    // To/from cam_core
    output logic [(`LANES)*(`ROW_BITS)-1:0]     rd_addr_lanes_flat,
    input  logic [(`LANES)*(`HASH_BITS)-1:0]    rd_hash_lanes_flat
`ifdef SIM_DEBUG
    ,output logic [(`NUM_ROWS)*(`SCORE_BITS)-1:0] score_debug_flat
`endif
);

    typedef enum logic [1:0] {  // state_t
        S_IDLE,  // Waiting for query_valid
        S_SCAN,  // Scanning rows in LANES-wide batches
        S_DONE   // Result ready, waiting for result_ready
    } state_t;

    state_t state_q, state_d;
    logic [(`HASH_BITS)-1:0] query_q;
    logic [(`ROW_BITS)-1:0] base_row_q, base_row_d;

    logic [(`ROW_BITS)-1:0] prev_best_idx [0:(`LANES)];
    logic [(`ROW_BITS)-1:0] next_best_idx [0:(`LANES)-1];
    logic [(`SCORE_BITS)-1:0] prev_best_score [0:(`LANES)];
    logic [(`SCORE_BITS)-1:0] next_best_score [0:(`LANES)-1];

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

    logic [(`ROW_BITS)-1:0] batch_best_idx;
    logic [(`SCORE_BITS)-1:0] batch_best_score;

    logic [(`ROW_BITS)-1:0] best_idx_q, best_idx_d;
    logic [(`SCORE_BITS)-1:0] best_score_q, best_score_d;

    assign query_ready  = (state_q == S_IDLE);
    assign result_valid = (state_q == S_DONE);
    assign result_row   = best_idx_q;
    assign result_score = best_score_q;
    assign busy         = (state_q == S_SCAN);

    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);

            assign rd_addr_lanes_flat[lane*`ROW_BITS +: `ROW_BITS] = lane_row[lane];
            assign row_hash = rd_hash_lanes_flat[lane*`HASH_BITS +: `HASH_BITS];

            assign effective_hash = row_hash;

            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_update (
                .best_index_i(prev_best_idx[lane]),
                .best_score_i(prev_best_score[lane]),
                .cand_index_i(lane_row[lane]),
                .cand_score_i(lane_valid[lane] ? lane_score[lane] : '0),
                .best_index_o(next_best_idx[lane]),
                .best_score_o(next_best_score[lane])
            );
        end
    endgenerate

    // Initialize chain with current batch seed
    assign prev_best_idx[0]   = best_idx_q;
    assign prev_best_score[0] = best_score_q;

    // Propagate per-lane results
    for (genvar l = 0; l < `LANES; l++) begin : chain_link
        assign prev_best_idx[l+1]   = next_best_idx[l];
        assign prev_best_score[l+1] = next_best_score[l];
    end

    assign batch_best_idx   = prev_best_idx[`LANES];
    assign batch_best_score = prev_best_score[`LANES];

    always_comb begin
        state_d      = state_q;
        base_row_d   = base_row_q;
        best_idx_d   = best_idx_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_idx_d   = TIE_BREAK_SENTINEL;    // Lower index wins tie-break
                    best_score_d = '0;
                end
            end

            S_SCAN: begin
                best_idx_d   = batch_best_idx;
                best_score_d = batch_best_score;

                if (base_row_q + `LANES >= `NUM_ROWS) begin
                    state_d = S_DONE;
                end else begin
                    base_row_d = base_row_q + `LANES;
                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_idx_q   <= '0;
            best_score_q <= '0;
        end else begin
            state_q      <= state_d;
            base_row_q   <= base_row_d;
            best_idx_q   <= best_idx_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

`ifdef SIM_DEBUG
    initial begin
        score_debug_flat = '0;
    end
`endif

endmodule