`timescale 1ns/1ps

// Testbench for jpeg_core.
//
// Usage (via vvp):
//   vvp sim.vvp +infile=image.jpg [+outfile=output.ppm]
//
// JPEG bytes are packed little-endian into 32-bit words:
//   data[7:0]   = first byte   (strb[0])
//   data[15:8]  = second byte  (strb[1])
//   data[23:16] = third byte   (strb[2])
//   data[31:24] = fourth byte  (strb[3])
//
// inport_last is intentionally left de-asserted; jpeg_input detects
// end-of-image via the JPEG EOI marker (0xFFD9) naturally.

module tb_jpeg_core;

// ---------------------------------------------------------------
// Tuneable limits
// ---------------------------------------------------------------
parameter MAX_FILE_BYTES = 1 * 1024 * 1024;  // 1 MB JPEG input
parameter MAX_IMG_PIXELS = 1024 * 1024;       // max output pixels
parameter TIMEOUT_CYCLES = 20_000_000;        // safety watchdog

// ---------------------------------------------------------------
// Clock / reset
// ---------------------------------------------------------------
reg clk = 0;
reg rst = 1;
always #5 clk = ~clk;   // 100 MHz

// ---------------------------------------------------------------
// DUT ports
// ---------------------------------------------------------------
reg         inport_valid = 0;
reg  [31:0] inport_data  = 0;
reg  [ 3:0] inport_strb  = 0;
reg         inport_last  = 0;
wire        inport_accept;

wire        outport_valid;
wire [15:0] outport_width;
wire [15:0] outport_height;
wire [15:0] outport_pixel_x;
wire [15:0] outport_pixel_y;
wire [ 7:0] outport_pixel_r;
wire [ 7:0] outport_pixel_g;
wire [ 7:0] outport_pixel_b;
reg         outport_accept = 1;

wire idle;

// ---------------------------------------------------------------
// DUT instance
// ---------------------------------------------------------------
jpeg_core #(.SUPPORT_WRITABLE_DHT(1)) dut (
    .clk_i            (clk),
    .rst_i            (rst),
    .inport_valid_i   (inport_valid),
    .inport_data_i    (inport_data),
    .inport_strb_i    (inport_strb),
    .inport_last_i    (inport_last),
    .inport_accept_o  (inport_accept),
    .outport_accept_i (outport_accept),
    .outport_valid_o  (outport_valid),
    .outport_width_o  (outport_width),
    .outport_height_o (outport_height),
    .outport_pixel_x_o(outport_pixel_x),
    .outport_pixel_y_o(outport_pixel_y),
    .outport_pixel_r_o(outport_pixel_r),
    .outport_pixel_g_o(outport_pixel_g),
    .outport_pixel_b_o(outport_pixel_b),
    .idle_o           (idle)
);

// ---------------------------------------------------------------
// Input file storage
// ---------------------------------------------------------------
reg [7:0] jpeg_mem [0:MAX_FILE_BYTES-1];
integer   jpeg_size;

// ---------------------------------------------------------------
// Output pixel storage  (flat: index = y*width + x)
// ---------------------------------------------------------------
reg [ 7:0] pix_r [0:MAX_IMG_PIXELS-1];
reg [ 7:0] pix_g [0:MAX_IMG_PIXELS-1];
reg [ 7:0] pix_b [0:MAX_IMG_PIXELS-1];
reg [15:0] img_width  = 0;
reg [15:0] img_height = 0;

// ---------------------------------------------------------------
// Capture pixels from DUT output
// ---------------------------------------------------------------
integer pix_idx;
always @(posedge clk) begin
    if (outport_valid && outport_accept) begin
        if (img_width == 0) begin
            img_width  <= outport_width;
            img_height <= outport_height;
        end
        pix_idx = outport_pixel_y * outport_width + outport_pixel_x;
        pix_r[pix_idx] <= outport_pixel_r;
        pix_g[pix_idx] <= outport_pixel_g;
        pix_b[pix_idx] <= outport_pixel_b;
    end
end

// ---------------------------------------------------------------
// Watchdog
// ---------------------------------------------------------------
integer cycle_count = 0;
always @(posedge clk) begin
    cycle_count <= cycle_count + 1;
    if (cycle_count >= TIMEOUT_CYCLES) begin
        $display("ERROR: simulation timeout after %0d cycles", TIMEOUT_CYCLES);
        $finish;
    end
end

// ---------------------------------------------------------------
// PPM writer task
// ---------------------------------------------------------------
reg [1023:0] out_filename;

task write_ppm;
    integer fd, x, y;
    begin
        fd = $fopen(out_filename, "wb");
        if (fd == 0) begin
            $display("ERROR: cannot open output file '%0s'", out_filename);
            $finish;
        end
        $fwrite(fd, "P6\n%0d %0d\n255\n", img_width, img_height);
        for (y = 0; y < img_height; y = y + 1)
            for (x = 0; x < img_width; x = x + 1)
                $fwrite(fd, "%c%c%c",
                    pix_r[y * img_width + x],
                    pix_g[y * img_width + x],
                    pix_b[y * img_width + x]);
        $fclose(fd);
    end
endtask

// ---------------------------------------------------------------
// Main stimulus
// ---------------------------------------------------------------
reg [1023:0] in_filename;
integer in_fd;
integer byte_idx;
integer remaining;
integer last_pixel_count;
integer quiescent_cycles;

integer pixel_count;
always @(posedge clk)
    if (outport_valid && outport_accept)
        pixel_count <= pixel_count + 1;
    else if (rst)
        pixel_count <= 0;

initial begin
    // --- argument parsing ---
    if (!$value$plusargs("infile=%s", in_filename)) begin
        $display("ERROR: +infile=<path> required");
        $finish;
    end
    if (!$value$plusargs("outfile=%s", out_filename))
        out_filename = "output.ppm";

    // --- load JPEG ---
    in_fd = $fopen(in_filename, "rb");
    if (in_fd == 0) begin
        $display("ERROR: cannot open '%0s'", in_filename);
        $finish;
    end
    jpeg_size = $fread(jpeg_mem, in_fd);
    $fclose(in_fd);
    $display("INFO: loaded %0d bytes from %0s", jpeg_size, in_filename);

    // --- reset ---
    repeat (8) @(posedge clk);
    @(negedge clk); rst = 0;
    repeat (2) @(posedge clk);

    // --- feed JPEG bytes as 32-bit little-endian words ---
    // Hold each word stable until inport_accept_o pulses, then advance.
    // inport_last is not used; EOI detection in jpeg_input handles termination.
    byte_idx = 0;
    while (byte_idx < jpeg_size) begin
        @(negedge clk);

        remaining = jpeg_size - byte_idx;
        inport_valid = 1;

        if (remaining >= 4) begin
            inport_data = { jpeg_mem[byte_idx+3],
                            jpeg_mem[byte_idx+2],
                            jpeg_mem[byte_idx+1],
                            jpeg_mem[byte_idx+0] };
            inport_strb = 4'hF;
        end else begin
            case (remaining)
                3: begin
                    inport_data = { 8'h00,
                                    jpeg_mem[byte_idx+2],
                                    jpeg_mem[byte_idx+1],
                                    jpeg_mem[byte_idx+0] };
                    inport_strb = 4'h7;
                end
                2: begin
                    inport_data = { 16'h0000,
                                    jpeg_mem[byte_idx+1],
                                    jpeg_mem[byte_idx+0] };
                    inport_strb = 4'h3;
                end
                default: begin   // 1
                    inport_data = { 24'h000000,
                                    jpeg_mem[byte_idx+0] };
                    inport_strb = 4'h1;
                end
            endcase
        end

        @(posedge clk);
        while (!inport_accept) @(posedge clk);

        byte_idx = byte_idx + ((remaining >= 4) ? 4 : remaining);
    end

    // De-assert after last word accepted
    @(negedge clk);
    inport_valid = 0;

    // --- wait for decoder to finish ---
    // Primary: wait for idle_o (requires EOF sentinel to drain pipeline).
    // Fallback: wait 5000 quiescent cycles after pixel output stops.
    begin : wait_idle
        quiescent_cycles = 0;
        last_pixel_count = -1;
        while (!idle) begin
            @(posedge clk);
            quiescent_cycles = quiescent_cycles + 1;
            if (pixel_count > 0 && pixel_count == last_pixel_count) begin
                if (quiescent_cycles > 5000) begin
                    disable wait_idle;
                end
            end else begin
                quiescent_cycles = 0;
                last_pixel_count = pixel_count;
            end
        end
    end
    repeat (20) @(posedge clk);

    // --- write result ---
    $display("INFO: decoded %0dx%0d image, writing to %0s",
             img_width, img_height, out_filename);
    write_ppm;
    $finish;
end

endmodule