`timescale 1ns / 1ps

module line_buffer #(
    parameter DATA_WIDTH=8,
    parameter IMAGE_SIZE=5,
    parameter NUM_LINES=3,
    parameter PADDING=1
)(
    input wire clk,
    input wire rst,
    input wire pixel_valid,
    input wire [DATA_WIDTH-1:0] pixel_in,
    output wire [DATA_WIDTH*NUM_LINES*NUM_LINES-1:0] window,
    output reg window_valid
);
    localparam IMAGE_SIZE_PADDED = IMAGE_SIZE+2*PADDING;
    localparam FIFO_DEPTH = IMAGE_SIZE * IMAGE_SIZE;

    // Padding registers
    reg [$clog2(IMAGE_SIZE_PADDED)-1:0] col_count_pad;
    reg [$clog2(IMAGE_SIZE_PADDED)-1:0] row_count_pad;
    reg [DATA_WIDTH-1:0] fifo [0:FIFO_DEPTH-1];
    reg [$clog2(FIFO_DEPTH)-1:0] write_ptr;
    reg [$clog2(FIFO_DEPTH)-1:0] read_ptr;
    reg fifo_empty;
    reg is_padding;
    
    // Padded pixel output
    reg [DATA_WIDTH-1:0] padded_pixel;
    reg padded_pixel_valid;

    // Line buffer registers
    reg [DATA_WIDTH-1:0] line_buf [0:NUM_LINES-1][0:IMAGE_SIZE_PADDED-1];
    reg [DATA_WIDTH-1:0] window_bufs [0:NUM_LINES-1][0:NUM_LINES-1];
    reg [$clog2(IMAGE_SIZE_PADDED)-1:0] buf_ptr;
    reg [$clog2(NUM_LINES)-1:0] row_count;
    reg [$clog2(IMAGE_SIZE*IMAGE_SIZE):0] window_count;
    reg frame_complete;

    integer i, j;

    // Padding logic
    always @(*) begin
        if (PADDING) begin
            is_padding = (row_count_pad == 0) || (row_count_pad == IMAGE_SIZE_PADDED-1) ||
                        (col_count_pad == 0) || (col_count_pad == IMAGE_SIZE_PADDED-1);
            
            if (!is_padding) begin
                padded_pixel = fifo[read_ptr];
                padded_pixel_valid = !fifo_empty;
            end else begin
                padded_pixel = 0;
                padded_pixel_valid = pixel_valid;
            end
        end else begin
            padded_pixel = pixel_in;
            padded_pixel_valid = pixel_valid;
        end
    end

    // FIFO and counter
    always @(posedge clk) begin
        if (rst) begin
            col_count_pad <= 0;
            row_count_pad <= 0;
            write_ptr <= 0;
            read_ptr <= 0;
            fifo_empty <= 1'b1;
        end else if (PADDING) begin
            if (pixel_valid) begin
                fifo[write_ptr] <= pixel_in;
                write_ptr <= (write_ptr == FIFO_DEPTH-1) ? 0 : write_ptr + 1;
                fifo_empty <= 1'b0;
            end

            if (!is_padding && !fifo_empty) begin
                read_ptr <= (read_ptr == FIFO_DEPTH-1) ? 0 : read_ptr + 1;
                fifo_empty <= (read_ptr + 1 == write_ptr);
            end

            if (col_count_pad == IMAGE_SIZE_PADDED-1) begin
                col_count_pad <= 0;
                row_count_pad <= (row_count_pad == IMAGE_SIZE_PADDED-1) ? 0 : row_count_pad + 1;
            end else begin
                col_count_pad <= col_count_pad + 1;
            end
        end
    end

    // Line buffer shifting
    always @(posedge clk) begin
        if (rst) begin
            buf_ptr <= 0;
            row_count <= 0;
            window_valid <= 0;
            window_count <= 0;
            frame_complete <= 0;
        end else begin
            for (i=NUM_LINES-1; i>0; i=i-1) begin
                line_buf[i][buf_ptr] <= line_buf[i-1][buf_ptr];
            end
            line_buf[0][buf_ptr] <= padded_pixel;

            if (buf_ptr == IMAGE_SIZE_PADDED-1) begin
                buf_ptr <= 0;
                row_count <= (row_count == NUM_LINES) ? row_count : row_count + 1;
            end else begin
                buf_ptr <= buf_ptr + 1;
            end
            
            window_valid <= (!frame_complete && row_count == NUM_LINES && buf_ptr >= NUM_LINES-1) ? 1 : 0;

            if (window_valid) begin
                if (window_count == (IMAGE_SIZE*IMAGE_SIZE)-1) begin
                    window_valid <= 0;
                    window_count <= 0;
                    frame_complete <= 1;
                end else begin
                    window_count <= window_count + 1;
                end
            end
        end
    end

    // Window buffer shifting
    always @(posedge clk) begin
        if (rst) begin
            for (i=0; i<NUM_LINES; i=i+1) begin
                for (j=0; j<NUM_LINES; j=j+1) begin
                    window_bufs[i][j] <= 0;
                end
            end
        end else begin
            for (i=0; i<NUM_LINES; i=i+1) begin
                window_bufs[i][NUM_LINES-1] <= line_buf[i][buf_ptr];
                for (j=NUM_LINES-1; j>0; j=j-1) begin
                    window_bufs[i][j-1] <= window_bufs[i][j];
                end
            end
        end
    end

    // Window output generation
    generate
        genvar k, l;
        for (k=0; k<NUM_LINES; k=k+1) begin
            for (l=0; l<NUM_LINES; l=l+1) begin
                assign window[((NUM_LINES-1-k)*NUM_LINES+l)*DATA_WIDTH+:DATA_WIDTH] = window_bufs[k][l];
            end
        end
    endgenerate

endmodule