`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,

    // Serial Top-K result interface
    output logic                              result_valid,
    input  logic                              result_ready,
    output logic [((`TOPK_K <= 1) ? 1 : $clog2(`TOPK_K))-1:0] result_rank,
    output logic [(`ROW_BITS)-1:0]            result_row,
    output logic [(`SCORE_BITS)-1:0]          result_score,
    output logic                              result_last,

    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
    ,output logic [($clog2(`NUM_ROWS+1))-1:0] consumed_candidates_debug
`endif
);
    //==========================================================================
    // State encoding
    //==========================================================================
    typedef enum logic [3:0] {
        S_IDLE,
        S_ISSUE_READ,
        S_WAIT_READ_RESP,
        S_WAIT_SCORE,
        S_STAGE_CANDIDATES,
        S_PUSH_CANDIDATES,
        S_DRAIN,
        S_SERIALIZE_RESULT,
        S_DONE
    } state_t;

    state_t state_q;

    //==========================================================================
    // Query / read tracking
    //==========================================================================
    logic [(`HASH_BITS)-1:0] query_q;
    logic [(`ROW_BITS)-1:0]  issue_base_q;
    logic [$clog2(`NUM_ROWS/`LANES+1)-1:0] issued_batches_q;
    logic [$clog2(`NUM_ROWS/`LANES+1)-1:0] returned_batches_q;
    logic [$clog2(`NUM_ROWS+1)-1:0]        pushed_candidates_q;
    logic [$clog2(`NUM_ROWS+1)-1:0]        consumed_candidates_q;
    // Total expected valid candidates across all batches
    logic [$clog2(`NUM_ROWS+1)-1:0]        expected_candidates_q;
    // Batch-level valid-lane mask for return tracking
    logic [(`LANES)-1:0]                   batch_lane_mask_q;

    //--------------------------------------------------------------------------
    // Read/noise staging registers (feed popcount pipeline)
    //--------------------------------------------------------------------------
    logic rd_stage_valid_q;
    logic [(`LANES)*(`ROW_BITS)-1:0]  rd_stage_row_ids_q;
    logic [(`LANES)*(`HASH_BITS)-1:0] rd_stage_hashes_q;
    logic [(`LANES)-1:0]              rd_stage_lane_valid_q;

    //--------------------------------------------------------------------------
    // Candidate staging registers (scored output from popcount, before FIFO)
    //--------------------------------------------------------------------------
    logic [(`LANES)-1:0]                    staged_valid_q;
    logic [(`LANES)*(`ROW_BITS)-1:0]        staged_rows_q;
    logic [(`LANES)*(`SCORE_BITS)-1:0]      staged_scores_q;
    logic [$clog2(`LANES)-1:0]              push_lane_idx_q;

    //==========================================================================
    // Popcount per-lane wires
    //==========================================================================
    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];

    //==========================================================================
    // Candidate FIFO interface
    //==========================================================================
    logic cand_valid;
    logic cand_ready;
    logic [(`ROW_BITS)-1:0] cand_row;
    logic [(`SCORE_BITS)-1:0] cand_score;
    wire cand_fire = cand_valid && cand_ready;

    //==========================================================================
    // FIFO wires
    //==========================================================================
    logic fifo_wr_valid;
    logic fifo_wr_ready;
    logic [(`ROW_BITS)-1:0] fifo_wr_row;
    logic [(`SCORE_BITS)-1:0] fifo_wr_score;
    logic fifo_rd_valid;
    logic fifo_rd_ready;
    logic [(`ROW_BITS)-1:0] fifo_rd_row;
    logic [(`SCORE_BITS)-1:0] fifo_rd_score;
    logic fifo_empty;

    //==========================================================================
    // Tracker wires
    //==========================================================================
    logic tracker_clear;
    logic tracker_candidate_ready;
    logic [(`TOPK_K)*(`ROW_BITS)-1:0]   tracker_topk_rows;
    logic [(`TOPK_K)*(`SCORE_BITS)-1:0] tracker_topk_scores;
    logic tracker_update_pending;

    //==========================================================================
    // Serializer wires
    //==========================================================================
    logic serializer_start;
    logic serializer_busy;
    logic serializer_done;
    logic serializer_result_valid;
    logic [((`TOPK_K <= 1) ? 1 : $clog2(`TOPK_K))-1:0] serializer_result_rank;
    logic [(`ROW_BITS)-1:0]   serializer_result_row;
    logic [(`SCORE_BITS)-1:0] serializer_result_score;
    logic serializer_result_last;

    //==========================================================================
    // Completion condition wires
    // Width-matched comparison targets to suppress Verilator WIDTHEXPAND.
    // The assigned constants fit exactly — WIDTHTRUNC is benign.
    /* verilator lint_off WIDTHTRUNC */
    wire [$clog2(`NUM_ROWS/`LANES+1)-1:0] num_batches_w = `NUM_ROWS / `LANES;
    wire [$clog2(`LANES)-1:0]             last_lane_w   = `LANES - 1;
    /* verilator lint_on WIDTHTRUNC */

    wire all_reads_issued      = (issued_batches_q     == num_batches_w);
    wire all_scores_returned   = (returned_batches_q   == num_batches_w);
    wire all_candidates_pushed = (pushed_candidates_q  == expected_candidates_q);
    wire all_candidates_consumed = (consumed_candidates_q == expected_candidates_q);
    wire scan_drained = all_reads_issued && all_scores_returned
                        && all_candidates_pushed
                        && all_candidates_consumed
                        && fifo_empty
                        && !tracker_update_pending;

    //==========================================================================
    // Combinational output assignments
    //==========================================================================
    assign query_ready    = (state_q == S_IDLE);
    assign rd_valid_o     = (state_q == S_ISSUE_READ);
    assign rd_base_row_o  = issue_base_q;
    assign busy           = (state_q != S_IDLE) || serializer_busy;

`ifdef SIM_DEBUG
    assign consumed_candidates_debug = consumed_candidates_q;
`endif

    // Result outputs come directly from the serializer
    assign result_valid = serializer_result_valid;
    assign result_rank  = serializer_result_rank;
    assign result_row   = serializer_result_row;
    assign result_score = serializer_result_score;
    assign result_last  = serializer_result_last;

    //==========================================================================
    // Generate: popcount pipeline per lane
    //==========================================================================
    generate
        for (genvar lane = 0; lane < `LANES; lane++) begin : gen_scores
            assign match_bits[lane] = ~(query_q ^ rd_stage_hashes_q[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_stage_valid_q && rd_stage_lane_valid_q[lane]),
                .row_i(rd_stage_row_ids_q[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

    //==========================================================================
    // Batch scores done: true when all valid lanes in the current batch
    // have their score_valid asserted, or vacuously true for an all-invalid batch.
    wire batch_scores_done;
    generate
        if (`LANES == 1) begin : gen_batch_done_1
            assign batch_scores_done = rd_stage_lane_valid_q[0]
                                       ? score_valid[0] : 1'b1;
        end else begin : gen_batch_done_n
            logic all_valid_ready;
            always_comb begin
                all_valid_ready = 1'b1;
                for (int l = 0; l < `LANES; l++) begin
                    if (rd_stage_lane_valid_q[l] && !score_valid[l]) begin
                        all_valid_ready = 1'b0;
                    end
                end
            end
            assign batch_scores_done = all_valid_ready;
        end
    endgenerate

    //==========================================================================
    // Candidate FIFO write interface
    //==========================================================================
    assign cand_valid = (state_q == S_PUSH_CANDIDATES)
                        && staged_valid_q[push_lane_idx_q];
    assign cand_row   = staged_rows_q[push_lane_idx_q * `ROW_BITS +: `ROW_BITS];
    assign cand_score = staged_scores_q[push_lane_idx_q * `SCORE_BITS +: `SCORE_BITS];
    assign cand_ready = fifo_wr_ready;

    assign fifo_wr_valid = cand_valid;
    assign fifo_wr_row   = cand_row;
    assign fifo_wr_score = cand_score;

    //==========================================================================
    // FIFO output
    //==========================================================================
    assign fifo_rd_ready = tracker_candidate_ready;

    wire fifo_rd_fire = fifo_rd_valid && fifo_rd_ready;

    //==========================================================================
    // Serializer start edge detection
    //==========================================================================
    logic start_ser_q1, start_ser_q2;
    wire  start_ser_raw = (state_q == S_DRAIN) && scan_drained;

    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            start_ser_q1 <= 1'b0;
            start_ser_q2 <= 1'b0;
        end else begin
            start_ser_q1 <= start_ser_raw;
            start_ser_q2 <= start_ser_q1;
        end
    end

    assign serializer_start = start_ser_q1 && !start_ser_q2;

    //==========================================================================
    // Tracker clear: pulse on new query
    //==========================================================================
    assign tracker_clear = (state_q == S_IDLE) && query_valid;

    //==========================================================================
    // Submodule instantiations
    //==========================================================================

    // --- candidate_fifo (full_o unused) ---
    /* verilator lint_off PINCONNECTEMPTY */
    candidate_fifo #(
        .DEPTH(`FIFO_DEPTH),
        .ROW_BITS(`ROW_BITS),
        .SCORE_BITS(`SCORE_BITS)
    ) u_candidate_fifo (
        .clk(clk),
        .rst_n(rst_n),
        .wr_valid_i(fifo_wr_valid),
        .wr_ready_o(fifo_wr_ready),
        .wr_row_i(fifo_wr_row),
        .wr_score_i(fifo_wr_score),
        .rd_valid_o(fifo_rd_valid),
        .rd_ready_i(fifo_rd_ready),
        .rd_row_o(fifo_rd_row),
        .rd_score_o(fifo_rd_score),
        .empty_o(fifo_empty),
        .full_o()              // intentionally unused
    );
    /* verilator lint_on PINCONNECTEMPTY */

    // --- topk_tracker ---
    topk_tracker #(
        .K(`TOPK_K),
        .ROW_BITS(`ROW_BITS),
        .SCORE_BITS(`SCORE_BITS)
    ) u_topk_tracker (
        .clk(clk),
        .rst_n(rst_n),
        .clear_i(tracker_clear),
        .candidate_valid_i(fifo_rd_valid),
        .candidate_ready_o(tracker_candidate_ready),
        .candidate_row_i(fifo_rd_row),
        .candidate_score_i(fifo_rd_score),
        .topk_rows_o(tracker_topk_rows),
        .topk_scores_o(tracker_topk_scores),
        .update_pending_o(tracker_update_pending)
    );

    // --- result_serializer ---
    result_serializer #(
        .K(`TOPK_K),
        .ROW_BITS(`ROW_BITS),
        .SCORE_BITS(`SCORE_BITS)
    ) u_result_serializer (
        .clk(clk),
        .rst_n(rst_n),
        .start_i(serializer_start),
        .busy_o(serializer_busy),
        .done_o(serializer_done),
        .topk_rows_i(tracker_topk_rows),
        .topk_scores_i(tracker_topk_scores),
        .result_valid_o(serializer_result_valid),
        .result_ready_i(result_ready),
        .result_rank_o(serializer_result_rank),
        .result_row_o(serializer_result_row),
        .result_score_o(serializer_result_score),
        .result_last_o(serializer_result_last)
    );

    //==========================================================================
    // Sequential: main FSM + registers
    //==========================================================================
    always_ff @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state_q              <= S_IDLE;
            query_q              <= '0;
            issue_base_q         <= '0;
            issued_batches_q     <= 0;
            returned_batches_q   <= 0;
            pushed_candidates_q  <= 0;
            consumed_candidates_q <= 0;
            expected_candidates_q <= 0;
            batch_lane_mask_q    <= '0;
            rd_stage_valid_q     <= 1'b0;
            rd_stage_row_ids_q   <= '0;
            rd_stage_hashes_q    <= '0;
            rd_stage_lane_valid_q <= '0;
            staged_valid_q       <= '0;
            staged_rows_q        <= '0;
            staged_scores_q      <= '0;
            push_lane_idx_q      <= 0;
`ifdef SIM_DEBUG
            score_debug_flat     <= '0;
`endif
        end else begin
            // Defaults for one-cycle flags
            rd_stage_valid_q <= 1'b0;

            // Consumed candidates counter (FIFO → Tracker handshake)
            if (fifo_rd_fire) begin
                consumed_candidates_q <= consumed_candidates_q + 1;
            end

            // State machine
            unique case (state_q)

                S_IDLE: begin
                    if (query_valid) begin
                        query_q              <= query_hash;
                        issue_base_q         <= '0;
                        issued_batches_q     <= 0;
                        returned_batches_q   <= 0;
                        pushed_candidates_q  <= 0;
                        consumed_candidates_q <= 0;
                        expected_candidates_q <= 0;
                        batch_lane_mask_q    <= '0;
                        push_lane_idx_q      <= 0;
                        staged_valid_q       <= '0;
`ifdef SIM_DEBUG
                        score_debug_flat     <= '0;
`endif
                        state_q <= S_ISSUE_READ;
                    end
                end

                S_ISSUE_READ: begin
                    // Emit rd_valid_o for exactly one cycle (combinational), then
                    // transition to S_WAIT_READ_RESP. This ensures the read-noise
                    // PRNG in cam_read_noise only advances once per batch issue.
                    state_q <= S_WAIT_READ_RESP;
                end

                S_WAIT_READ_RESP: begin
                    if (rd_valid_i) begin
                        // Capture staging data for popcount pipeline
                        rd_stage_valid_q      <= 1'b1;
                        rd_stage_row_ids_q    <= rd_row_ids_i;
                        rd_stage_hashes_q     <= rd_hashes_i;
                        rd_stage_lane_valid_q <= rd_lane_valid_i;
                        // Remember the lane mask for this batch (used in S_WAIT_SCORE
                        // to determine batch_scores_done).
                        batch_lane_mask_q     <= rd_lane_valid_i;
                        issue_base_q          <= issue_base_q + `LANES;
                        issued_batches_q      <= issued_batches_q + 1;
                        // Accumulate total expected valid candidates
                        expected_candidates_q <= expected_candidates_q
                                                + $countones(rd_lane_valid_i);
                        state_q <= S_WAIT_SCORE;
                    end
                end

                S_WAIT_SCORE: begin
                    // rd_stage_valid_q already deasserted by default above
                    if (batch_scores_done) begin
                        returned_batches_q <= returned_batches_q + 1;

                        // Capture per-lane scores and rows into staging registers
                        for (int lane = 0; lane < `LANES; lane++) begin
                            if (batch_lane_mask_q[lane]) begin
                                staged_valid_q[lane] <= 1'b1;
                                staged_rows_q[lane*`ROW_BITS +: `ROW_BITS] <= score_row[lane];
                                staged_scores_q[lane*`SCORE_BITS +: `SCORE_BITS] <= lane_score[lane];
                            end else begin
                                staged_valid_q[lane] <= 1'b0;
                            end
                        end

`ifdef SIM_DEBUG
                        for (int lane = 0; lane < `LANES; lane++) begin
                            if (batch_lane_mask_q[lane]) begin
                                score_debug_flat[score_row[lane]*`SCORE_BITS +: `SCORE_BITS]
                                    <= lane_score[lane];
                            end
                        end
`endif
                        state_q <= S_STAGE_CANDIDATES;
                    end
                end

                S_STAGE_CANDIDATES: begin
                    state_q <= S_PUSH_CANDIDATES;
                end

                S_PUSH_CANDIDATES: begin
                    // Advance if the current lane fired or wasn't valid
                    if (cand_fire || !staged_valid_q[push_lane_idx_q]) begin
                        if (cand_fire) begin
                            pushed_candidates_q <= pushed_candidates_q + 1;
                        end

                        if (push_lane_idx_q == last_lane_w) begin
                            // Last lane processed — decide next state
                            push_lane_idx_q <= 0;
                            if (all_reads_issued) begin
                                state_q <= S_DRAIN;
                            end else begin
                                state_q <= S_ISSUE_READ;
                            end
                        end else begin
                            push_lane_idx_q <= push_lane_idx_q + 1;
                        end
                    end
                end

                S_DRAIN: begin
                    if (scan_drained) begin
                        state_q <= S_SERIALIZE_RESULT;
                    end
                end

                S_SERIALIZE_RESULT: begin
                    if (serializer_done) begin
                        state_q <= S_DONE;
                    end
                end

                S_DONE: begin
                    state_q <= S_IDLE;
                end

                default: state_q <= S_IDLE;
            endcase
        end
    end

endmodule