`include "gfx/gfx_defs.sv"

module gfx_sp_fetch
(
	input  logic          clk,
	                      rst_n,

	input  logic          fetch_waitrequest,
	                      fetch_readdatavalid,
	input  vram_word      fetch_readdata,
	output vram_addr      fetch_address,
	output logic          fetch_read,

	input  logic          program_start,
	input  cmd_insn_ptr   program_header_base,
	input  cmd_word       program_header_size,
	output logic          running,

	input  logic          batch_end,
	output vram_insn_addr batch_base,
	output logic          batch_start,
	output cmd_word       batch_length,

	input  logic          ready,
	output logic          valid,
	output insn_word      insn,
	output logic          clear_lanes
);

	localparam ENTRY_SIZE = 4;

	logic break_loop, entry_end, fifo_down_safe, fifo_empty, fifo_put,
	      header_continue, insn_read, insn_readdatavalid, insn_waitrequest;

	cmd_word header_count;
	insn_word code_length, code_read_ptr, code_fetch_ptr, insn_readdata, entry_data[ENTRY_SIZE];
	vram_insn_addr code_base, insn_address, header_ptr;
	logic[$clog2(ENTRY_SIZE - 1):0] entry_fetch_count, entry_read_count;

	enum int unsigned
	{
		IDLE,
		HEADER,
		LOOP,
		FLUSH
	} state;

	struct packed
	{
		insn_word insn;
		logic     clear_lanes;
	} fifo_in, fifo_out;

	assign insn = fifo_out.insn;
	assign clear_lanes = fifo_out.clear_lanes;

	assign entry_end = entry_read_count == ENTRY_SIZE - 1;
	assign header_continue = header_count != 0;
	assign break_loop = batch_end && (!insn_read || !insn_waitrequest);

	function vram_insn_addr base_from_word(insn_word in);
		base_from_word = in[`GFX_INSN_SUBWORD_BITS +: $bits(vram_insn_addr)];
	endfunction

	assign code_base = base_from_word(entry_data[0]);
	assign batch_base = base_from_word(entry_data[2]);
	assign code_length = entry_data[1];
	assign batch_length = entry_data[3];

	gfx_sp_widener #(.WIDTH($bits(vram_insn_addr))) insn_bus
	(
		.wide_read(insn_read),
		.wide_address(insn_address),
		.wide_readdata(insn_readdata),
		.wide_waitrequest(insn_waitrequest),
		.wide_readdatavalid(insn_readdatavalid),
		.word_read(fetch_read),
		.word_address(fetch_address),
		.word_readdata(fetch_readdata),
		.word_waitrequest(fetch_waitrequest),
		.word_readdatavalid(fetch_readdatavalid),
		.*
	);

	gfx_fifo #(.WIDTH($bits(fifo_in)), .DEPTH(`GFX_FETCH_FIFO_DEPTH)) insn_fifo
	(
		.in(fifo_in),
		.out(fifo_out),
		.in_ready(),
		.in_valid(fifo_put),
		.out_ready(ready),
		.out_valid(valid),
		.*
	);

	gfx_fifo_overflow #(.DEPTH(`GFX_FETCH_FIFO_DEPTH)) overflow
	(
		.down(insn_read && !insn_waitrequest),
		.empty(fifo_empty),
		.down_safe(fifo_down_safe),
		.out_ready(ready),
		.out_valid(valid),
		.*
	);

	always_ff @(posedge clk or negedge rst_n)
		if (!rst_n) begin
			state <= IDLE;
			running <= 0;
			fifo_put <= 0;
			insn_read <= 0;
			batch_start <= 0;
		end else unique case (state)
			IDLE:
				if (program_start) begin
					state <= HEADER;
					running <= 1;
					insn_read <= 1;
				end

			HEADER: begin
				if (insn_read && !insn_waitrequest)
					insn_read <= entry_fetch_count != ENTRY_SIZE - 1;

				if (insn_readdatavalid && entry_end) begin
					state <= LOOP;
					insn_read <= 1;
					batch_start <= 1;
				end
			end

			LOOP: begin
				fifo_put <= 0;
				batch_start <= 0;

				if (!insn_read || !insn_waitrequest)
					insn_read <= fifo_down_safe;

				if (break_loop) begin
					state <= FLUSH;
					insn_read <= 0;
				end

				if (insn_readdatavalid)
					fifo_put <= 1;
			end

			FLUSH: begin
				fifo_put <= 0;

				if (fifo_empty) begin
					state <= header_continue ? HEADER : IDLE;
					running <= header_continue;
					insn_read <= header_continue;
				end
			end
		endcase

	always_ff @(posedge clk)
		unique case (state)
			IDLE:
				if (program_start) begin
					header_ptr <= program_header_base.addr;
					header_count <= program_header_size;
					insn_address <= program_header_base.addr;

					entry_read_count <= 0;
					entry_fetch_count <= 0;
				end

			HEADER: begin
				code_read_ptr <= 0;
				code_fetch_ptr <= 0;

				if (!insn_waitrequest) begin
					insn_address <= insn_address + 1;
					entry_fetch_count <= entry_fetch_count + 1;
				end

				if (insn_read && !insn_waitrequest)
					header_ptr <= header_ptr + 1;

				if (insn_readdatavalid) begin
					entry_read_count <= entry_read_count + 1;

					for (integer i = 0; i < ENTRY_SIZE - 1; ++i)
						entry_data[i] <= entry_data[i + 1];

					entry_data[ENTRY_SIZE - 1] <= insn_readdata;

					if (entry_end)
						insn_address <= base_from_word(entry_data[1]);
				end
			end

			LOOP: begin
				if (insn_read && !insn_waitrequest) begin
					insn_address <= insn_address + 1;
					code_fetch_ptr <= code_fetch_ptr + 1;

					if (code_fetch_ptr == code_length) begin
						insn_address <= code_base;
						code_fetch_ptr <= 0;
					end
				end

				if (insn_readdatavalid) begin
					fifo_in.insn <= insn_readdata;
					fifo_in.clear_lanes <= code_read_ptr == 0;

					code_read_ptr <= code_read_ptr + 1;
					if (code_read_ptr == code_length)
						code_read_ptr <= 0;
				end
			end

			FLUSH:
				if (fifo_empty) begin
					header_count <= header_count - 1;
					insn_address <= header_ptr;
				end
		endcase

endmodule