module poly1305_friendly_modular_mult #(
    parameter DATA_WIDTH = 128,
    parameter ACC_WIDTH = 130
) (
    input  logic                    i_clk,
    input  logic                    i_rst,

    input  logic                    i_valid,
    output logic                    o_ready,
    input  logic [DATA_WIDTH-1:0]   i_data,
    input  logic [ACC_WIDTH-1:0]    i_accumulator,

    output logic                    o_valid,
    output logic [ACC_WIDTH-1:0]    o_result
);

localparam [129:0] PRIME = (1 << 130) - 5;

logic [2:0] state_counter, state_counter_next;

logic [2:0] state_counter_p [4];

logic [ACC_WIDTH-1:0] accumulator, accumulator_next;    // accumulator is outgoing

logic [DATA_WIDTH-1:0] data, data_next;
logic [ACC_WIDTH-1:0] h, h_next;    // h is incoming

logic [DATA_WIDTH+26-1:0] mult_product, mult_product_next;

logic [ACC_WIDTH-1:0] modulo_result;

assign o_ready = state_counter >= 3'h4;
assign o_result = accumulator;

always_ff @(posedge i_clk) begin
    if (i_rst) begin
        state_counter <= 3'h5;
        state_counter_p <= '{default: 3'h5};
    end else begin
        state_counter <= state_counter_next;
        accumulator <= accumulator_next;
        data <= data_next;
        h <= h_next;
        mult_product <= mult_product_next;

        state_counter_p[0] <= state_counter;

        o_valid <= state_counter_p[3] == 3'h4;

        for (int i = 1; i < 4; i++) begin
            state_counter_p[i] <= state_counter_p[i-1];
        end
    end
end

always_comb begin
    data_next = data;
    h_next = h;

    state_counter_next = state_counter;

    accumulator_next = '0;
    mult_product_next = '0;


    if (state_counter >= 3'h4 && i_valid) begin
        data_next = i_data;
        h_next = i_accumulator;
        state_counter_next = '0;
    end

    if (state_counter < 3'h5) begin
        mult_product_next = h[state_counter*26 +: 26] * data;
        state_counter_next = state_counter + 1;
    end

    if (state_counter_p[3] == '0) begin
        accumulator_next = modulo_result;
    end else begin
        if (accumulator + modulo_result > PRIME) begin
            accumulator_next = accumulator + modulo_result - PRIME;
        end else begin
            accumulator_next = accumulator + modulo_result;
        end
    end

end

poly1305_friendly_modulo u_modulo (
    .i_clk          (i_clk),
    .i_rst          (i_rst),

    .i_valid        ('1),
    .i_val          ((2*ACC_WIDTH)'(mult_product)),
    .i_shift_amount (state_counter_p[0]),

    .o_valid        (),
    .o_result       (modulo_result)
);

endmodule