`ifndef ASS_3_SV
`define ASS_3_SV

`include "rom.sv"
`include "ram.sv"
`include "register_file.sv"
`include "instruction_counter.sv"
`include "control_unit.sv"
`include "alu.sv"
`include "shifter.sv"
`include "sign_extender.sv"
`include "gatereg.sv"
`include "hazard_unit.sv"

module assignment3 #(parameter progsrc = "", parameter width=5, parameter romsz=9, parameter ramsz=10)(
    input logic [0:0] clk, input logic [0:0] rst
);	
    logic [31:0] regs_out [32];

	// Inputs
	logic [0:0] hazard_m2r_0, hazard_m2r_1, hazard_m2r_2;
	logic [0:0] hazard_reg_we_1, hazard_reg_we_2, hazard_reg_we_3;
	logic [0:0] hazard_branch_taken, hazard_jump_taken;
	logic [4:0] hazard_rt_0, hazard_rt_1;
	logic [4:0] hazard_rs_0, hazard_rs_1;
	logic [4:0] hazard_write_reg_0, hazard_write_reg_1, hazard_write_reg_2;
	// Outputs
	logic [1:0] hazard_forwardAE, hazard_forwardBE;
	logic [0:0] hazard_forwardAD, hazard_forwardBD;
	logic [0:0] hazard_stallF, hazard_stallD;
	logic [0:0] hazard_flushE, hazard_flushD;
	hazard_unit hazard_unit(
		hazard_m2r_1, hazard_rt_1, hazard_rs_1,
		hazard_reg_we_1, hazard_write_reg_0, /*>Missing*/ hazard_m2r_1,
		hazard_reg_we_2, hazard_write_reg_1, /*>Missing*/ hazard_m2r_2,
		hazard_reg_we_3, hazard_write_reg_2, // MISSING,
		hazard_m2r_0 /* MISSING */, hazard_rs_0, hazard_rt_0,
		hazard_branch_taken, hazard_jump_taken,	// NOTE: We aren't doing branch prediction. Can just be control unit branch
		hazard_forwardAE, hazard_forwardBE,
		hazard_forwardAD, hazard_forwardBD,	// MISSING
		hazard_stallF, hazard_stallD,
		hazard_flushE, hazard_flushD
	);

	logic [2**width-1:0] pc_addr, pc_addr_1, pc_addr_2;
	logic [2**width-1:0] pc_addr_write;
	logic [0:0] pc_we;
	instruction_counter #(width) PC (clk, rst, ~hazard_stallF, pc_we, pc_addr_write, pc_addr);
	gatereg #(width, 1) gate_pc_addr(clk, rst, ~hazard_stallF, pc_addr, pc_addr_1); // TODO: Verify if this needs to exist
	gatereg #(width, 1) gate_pc_addr_1(clk, rst | (hazard_branch_taken | hazard_jump_taken), ~hazard_stallD, pc_addr_1, pc_addr_2);


	logic [2**width-1:0] rom_rd;
	logic [2**width-1:0] instruction, instruction_1;
	rom #(progsrc, romsz, width) ROM (pc_addr_1[romsz-1:0], rom_rd);
	gatereg #(width, 1) gate_instruction(clk, rst | (hazard_branch_taken | hazard_jump_taken), ~hazard_stallD, rom_rd, instruction);
	gatereg #(width, 1) gate_instruction_1(clk, rst | hazard_flushE, 1'b1, instruction, instruction_1);


	logic [0:0] cu_m2r, cu_m2w, cu_alu_src, cu_rdst, cu_rw;
	logic [0:0] cu_m2w_1, cu_rw_1, cu_m2r_1;
	logic [0:0] cu_m2r_2, cu_rw_2;
	logic [3:0] cu_alu_fn, ALU_FN;
	logic [0:0] B, MEM_WE, RDST, ALU_SRC, REG_WE, REGW_MEM_OR_RD;
	control_unit CONTROL_UNIT (instruction[31:26], instruction[5:0],
		cu_m2r, cu_m2w, B, cu_alu_fn, cu_alu_src, cu_rdst, cu_rw
	);
	gatereg #(0, 1) gate_cu_rdst (clk, rst | hazard_flushE, 1'b1, cu_rdst, RDST);
	gatereg #(0, 1) gate_cu_alu_src (clk, rst | hazard_flushE, 1'b1, cu_alu_src, ALU_SRC);
	gatereg #(2, 1) gate_cu_alu_fn (clk, rst | hazard_flushE, 1'b1, cu_alu_fn, ALU_FN);
	gatereg #(0, 1) gate_cu_m2w(clk, rst | hazard_flushE, 1'b1, cu_m2w, cu_m2w_1);
	gatereg #(0, 1) gate_cu_rw(clk, rst | hazard_flushE, 1'b1, cu_rw, cu_rw_1);
	gatereg #(0, 1) gate_cu_m2r(clk, rst | hazard_flushE, 1'b1, cu_m2r, cu_m2r_1);
	
	gatereg #(0, 1) gate_cu_m2r_1 (clk, rst, 1'b1, cu_m2r_1, cu_m2r_2);
	gatereg #(0, 1) gate_cu_m2r_2 (clk, rst, 1'b1, cu_m2r_2, REGW_MEM_OR_RD);

	gatereg #(0, 1) gate_cu_m2w_1 (clk, rst, 1'b1, cu_m2w_1, MEM_WE);
	
	gatereg #(0, 1) gate_cu_rw_1 (clk, rst, 1'b1, cu_rw_1, cu_rw_2);
	gatereg #(0, 1) gate_cu_rw_2 (clk, rst, 1'b1, cu_rw_2, REG_WE);
	
	assign hazard_m2r_0 = cu_m2r_1;
	assign hazard_m2r_1 = cu_m2r_2;
	assign hazard_m2r_2 = REGW_MEM_OR_RD;
	assign hazard_reg_we_1 = cu_rw_1;
	assign hazard_reg_we_2 = cu_rw_2;
	assign hazard_reg_we_3 = REG_WE;

	logic [2**width-1:0] regf_rs, regf_rs_1;
	logic [2**width-1:0] regf_rt, regf_rt_1, regf_rt_2;
	logic [7:0] regf_a3, regf_a3_1, regf_a3_2;
	logic [2**width-1:0] rd, rd_mem, rd_alu;
	assign rd = (REGW_MEM_OR_RD) ? rd_mem : rd_alu;
	assign regf_a3 = (RDST) ? instruction_1[15:11] : instruction_1[20:16];
	gatereg #(3, 1) gate_regf_a3 (clk, rst, 1'b1, regf_a3, regf_a3_1);
	gatereg #(3, 1) gate_regf_a3_1 (clk, rst, 1'b1, regf_a3_1, regf_a3_2);
	register_file #(width, 5) REGS (
		~clk, REG_WE, 
		regf_a3_2[4:0], rd,
		instruction[25:21], regf_rs,
		instruction[20:16], regf_rt,
		regs_out
	);
	gatereg #(width, 1) gate_regf_rs(clk, rst | hazard_flushE, 1'b1, regf_rs, regf_rs_1);
	gatereg #(width, 1) gate_regf_rt(clk, rst | hazard_flushE, 1'b1, regf_rt, regf_rt_1);
	assign hazard_write_reg_0 = regf_a3;
	assign hazard_write_reg_1 = regf_a3_1;
	assign hazard_write_reg_2 = regf_a3_2;
	assign hazard_rs_0 = instruction[25:21];
	assign hazard_rt_0 = instruction[20:16];
	assign hazard_rs_1 = instruction_1[25:21];
	assign hazard_rt_1 = instruction_1[20:16];

	gatereg #(width, 1) gate_regf_rt_1(clk, rst, 1'b1, regf_rt_1, regf_rt_2); // TODO: Remove - is WriteDataE, not rt
	
	logic [2**width-1:0] sign_ext_imm, rt_imm;
	sign_extender #(4, width) SIGNEXT(1'b1, instruction[15:0], sign_ext_imm);
	gatereg #(width, 1) gate_sign_ext_imm(clk, rst, 1'b1, sign_ext_imm, rt_imm);

	// Inputs
	logic [2**width-1:0] alu_in_a, alu_in_b;
	always_comb begin
		case (hazard_forwardAE)
			2'b00: alu_in_a = regf_rs_1;
			2'b01: alu_in_a = rd;
			2'b10: alu_in_a = alu_out_1;
			default: alu_in_a = regf_rs_1;
		endcase
		case (hazard_forwardBE)
			2'b00: alu_in_b = regf_rt_1;
			2'b01: alu_in_b = rd;
			2'b10: alu_in_b = alu_out_1;
			default: alu_in_b = regf_rt_1;
		endcase
	end
	// Outputs
	logic [2**width-1:0] alu_out, alu_out_1, alu_out_2;
	logic [3:0] alu_flags, alu_flags_1;
	alu #(width) ALU (ALU_FN, alu_in_a, (ALU_SRC) ? rt_imm : alu_in_b, alu_out, alu_flags[2:0]);
	gatereg #(width, 1) gate_alu_out (clk, rst, 1'b1, alu_out, alu_out_1);
	gatereg #(width, 1) gate_alu_out_1 (clk, rst, 1'b1, alu_out_1, alu_out_2);
	gatereg #(2, 1) gate_alu_flags (clk, rst, 1'b1, alu_flags, alu_flags_1);
	assign rd_alu = alu_out_2;

	logic [2**width-1:0] new_pc_addr, new_pc_addr_1;
	assign new_pc_addr = (sign_ext_imm << 2) + pc_addr_2;
	gatereg #(width, 1) gate_new_pc_addr (clk, rst, 1'b1, new_pc_addr, new_pc_addr_1);
	

	assign pc_addr_write = new_pc_addr_1;
	assign pc_we = B & ((hazard_forwardAD ? alu_out_1 : regf_rs) == (hazard_forwardBD ? alu_out_1 : regf_rt));
	assign hazard_branch_taken = pc_we;
	assign hazard_jump_taken = pc_we;
	logic [2**width-1:0] mem_out, mem_out_1;
	ram #(ramsz, width) RAM (clk, rst, alu_out_1[ramsz-1:0], MEM_WE, alu_in_b, mem_out);
	gatereg #(width, 1) gate_mem_out (clk, rst, 1'b1, mem_out, mem_out_1);
	assign rd_mem = mem_out_1;
endmodule
`endif