import ProcTypes;

import RegisterFile;
import Decode;
import Execute;

// MainMemory includes the MainMemory port
// and the SingleCycleMemory modules
import MainMemory;

// Milestone 2: Select which cache(s) to import
import DirectMappedCache;
//import TwoWayCache;


typedef struct {
    // Add anything you need
    // NOTE: If you leave F2D empty, Maybe#(F2D) f2d will just be the valid bit.
    // In this case, you can set f2d to valid by writing: f2d <= Valid(?);
} F2D;

typedef struct {
    Word pc;
    DecodedInst dInst;
    Word rVal1;
    Word rVal2;
    // Add anything you need
} D2E;

typedef struct {
    Word pc;
    IType iType;
    Maybe#(RIndx) dst;
    Word data;
    // Add anything you need
} E2W;

module Processor;
    Reg#(Word) pc(0);
    RegisterFile rf;

    // Uncomment and use these in the processor
    // rule for useful debugging help tracking
    // cycles and instructions completed
    //Reg#(Bit#(64)) cycles(0);
    //Reg#(Bit#(64)) insts(0);

    // Single Cycle Memories
    // Milestone 2: comment these declarations out
    //SingleCycleMemory iMem;
    //SingleCycleMemory dMem;

    // Line-based Main Memories to use with caches
    // Milestone 2: uncomment these declarations
    MainMemory iMem;
    MainMemory dMem;

    // Milestone 2: choose which type of cache to use
    // NOTE: The CacheWrapper#() type of the dCache
    // is critical, as it is used for detecting
    // that the program has finished running
    // iCache: Direct-Mapped or Two-Way Set Associative
    DirectMappedCache iCache(iMem);
    //TwoWayCache iCache(iMem);
    // dCache: Direct-Mapped or Two-Way Set Associative
    CacheWrapper#(DirectMappedCache) dCache(dMem);
    //CacheWrapper#(TwoWayCache) dCache(dMem);

    // Pipeline Registers
    Reg#(Maybe#(F2D)) f2d(Invalid);
    Reg#(Maybe#(D2E)) d2e(Invalid);
    Reg#(Maybe#(E2W)) e2w(Invalid);

    Reg#(Bool) reject_if(False);
    Reg#(Bool) repeat_if(False);
    // Performance counters
    // These are used for debugging, but are not needed for correctness; you
    // can modify or remove them
    Reg#(Bit#(32)) cycle(0);
    Reg#(Bit#(32)) instrs(0);

    rule tick;
        //////////////////////
        // Debugging Helper //
        //////////////////////
        cycle <= cycle + 1;
        // Uncomment the code below to print pipeline state each cycle
        begin
            $write("[%d] F|D|E|W PCs: ", cycle);
            $write("%x|", pc);
            if (isValid(f2d)) $write("%x|", pc); else $write("%8s|", "NOP");  // F and D share pc
            if (isValid(d2e)) $write("%x|", fromMaybe(?, d2e).pc); else $write("%8s|", "NOP");
            if (isValid(e2w)) $write("%x", fromMaybe(?, e2w).pc); else $write("%8s", "NOP");
            $display("  instrs: %d", instrs);
        end

        /////////////////////
        // Writeback Stage //
        /////////////////////
        // Signals used by decode to handle data hazards
        Maybe#(RIndx) dstW = Invalid;
        Word dataW = ?;

        // Milestone 2:
        // Signal for dCache-induced stalls
        Bool dStall = False;

        if (isValid(e2w)) begin
            let e2w_v = fromMaybe(?, e2w);

            //////////////////////////
            // TODO: Your code here //
            //////////////////////////
            if (e2w_v.iType == LOAD) begin
		if (isValid(dCache.data)) begin
                	rf.wr = Valid(RegWriteArgs{index: fromMaybe(?, e2w_v.dst), data: fromMaybe(?, dCache.data)}); // Load from data memory (dMem) if the instruction requires it
                	dstW = Valid(fromMaybe(?, e2w_v.dst));
                	$display("loading into register r", fromMaybe(?, e2w_v.dst));
			$display("loading data from mem: ", fromMaybe(?, dMem.data));
		end else begin
			if (!(dCache.reqEnabled)) begin
				dStall = True;
			end
			rf.wr = Invalid;
			dstW = Invalid;
		end

            end
            else if (isValid(e2w_v.dst)) begin
                dstW = Valid(fromMaybe(?, e2w_v.dst));
                rf.wr = Valid(RegWriteArgs{index: fromMaybe(?, e2w_v.dst), data: e2w_v.data}); // Write to a register if the instruction requires it
		$display("loading into register r", fromMaybe(?, e2w_v.dst));
                $display("loading data from exe: ", e2w_v.data);
            end else begin
                rf.wr = Invalid;
            end

            // Count every instruction that leaves the writeback stage
            // (for debugging help)
            instrs <= instrs + 1;

            // BEGIN: DO NOT REMOVE OR MODIFY THIS CODE
            // If unsupported instruction, stops simulation and print the state of the processor
            if (e2w_v.iType == Unsupported) begin
                $display("Reached unsupported instruction");
                $display("Dumping the state of the processor");
                $display("pc = 0x%x", e2w_v.pc);
                $display(rf.fshow);
                $display("Quitting simulation.");
                $finish;
            end
            // END: DO NOT REMOVE OR MODIFY THIS CODE

        end else begin
            rf.wr = Invalid;
        end

        ///////////////////
        // Execute Stage //
        ///////////////////

        // Signals used to handle mispredictions
        Bool annul = False;
        Word redirectPC = ?;

        // Signals used by decode to handle data hazards
        Maybe#(RIndx) dstE = Invalid;
        Word dataE = ?;

        if (isValid(d2e)) begin
            let d2e_v = fromMaybe(?, d2e);
	    if (dStall) begin
	    	e2w <= e2w;

	    end else begin

            	ExecInst eInst = execute(d2e_v.dInst, d2e_v.rVal1, d2e_v.rVal2, d2e_v.pc);
            	if (eInst.iType == LOAD) begin
			if (dCache.reqEnabled) begin
                		dCache.req = Valid(MemReq{addr: eInst.addr, op: Ld, data: ?});
                		//eInst.data = dMem.data;
			end else begin
				dStall = True;
				dCache.req = Invalid;
			end
            	end else if (eInst.iType == STORE) begin
                	if (dCache.reqEnabled) begin
				dCache.req = Valid(MemReq{addr: eInst.addr, data: eInst.data, op: St}); // ** FIX THIS LINE
            		end else begin
				dStall = True;
				dCache.req = Invalid;
			end
		end
		else begin
			dCache.req = Invalid;
		end
		if (dStall) begin
            		E2W e_to_writeback = E2W{pc: d2e_v.pc, iType: eInst.iType, dst: eInst.dst, data: eInst.data};
            		e2w <= Valid(e_to_writeback);
		end
		else begin
			e2w <= Invalid;
		end
            	if (isValid(eInst.dst)) begin
                	dstE = Valid(fromMaybe(?, eInst.dst));
            	end

            	if (eInst.nextPc != d2e_v.pc + 4) begin
                	annul = True;
                	redirectPC = eInst.nextPc;
            	end
	   end

        end else begin
            e2w <= Invalid;
        end


        //////////////////
        // Decode Stage //
        //////////////////

        // Signal for decode stalls
        Bool stall = False;

        if (isValid(f2d)) begin
	    if (dStall) begin
		d2e <= d2e;
	    end
	    else begin
            	if (!(isValid(iCache.data)) && !(iCache.reqEnabled)) begin // (in middle of a) CACHE MISS, stall
                	stall = True;
                	d2e <= Invalid;
                	if (annul) begin
                    	reject_if <= True;
                	end
            	end
	    	else if (!(isValid(iCache.data))) begin
			d2e <= Invalid; // stall decode but not fetch
	    	end
            	else begin
                	DecodedInst dInst = decode(fromMaybe(?, iCache.data));
                	// *** a MAYBE of type LINE not A WORD.

                	// Read source registers
                	Word rVal1 = rf.rd1(dInst.src1);
                	Word rVal2 = rf.rd2(dInst.src2);

                	if (isValid(dstE) && (fromMaybe(?, dstE) == dInst.src1 || fromMaybe(?, dstE) == dInst.src2)) begin
                    		stall = True;
	 	    		if (!(iCache.reqEnabled)) begin
		    			repeat_if <= True;
		    		end
		    		$display("e_hazard");
                	end
                	else if (isValid(dstW) && (fromMaybe(?, dstW) == dInst.src1 || fromMaybe(?, dstW) == dInst.src2)) begin
                    		stall = True;
		    		$display("w_hazard");
		    		if (!(iCache.reqEnabled)) begin
                        		repeat_if <= True;
                    		end
                	end

                	let f2d_v = fromMaybe(?, f2d);

                	D2E d_to_execute = D2E{pc: pc, dInst: dInst, rVal1: rVal1, rVal2: rVal2};

                	if (reject_if) begin
                    		d2e <= Invalid;
                    		reject_if <= False;
                	end

                	else if (stall || annul) begin
                    		d2e <= Invalid;
                	end

                	else begin
                    		d2e <= Valid(d_to_execute);
                	end
            	end

	    end
        end else begin
            d2e <= Invalid;
        end

        /////////////////
        // Fetch Stage //
        /////////////////


        if (isValid(f2d)) begin //aka if not first instr (more accurate: first time doing fetch code)
            let nextpc = pc;
            if (annul) begin
                nextpc = redirectPC;
            end
	    else if (stall || dStall) begin
	    	nextpc = pc;
	    end
            else if (iCache.reqEnabled) begin
        	if (repeat_if) begin
			nextpc = pc;
		end
		else begin
			nextpc = pc + 4;
		end
            end
            // else begin
            //     nextpc = pc;
            // end
            if (iCache.reqEnabled) begin
		if (repeat_if) begin
			repeat_if <= False;
		end
	    	iCache.req = Valid(MemReq{addr: nextpc, op: Ld, data: ?}); // do regardless of stall.
            end
	    pc <= nextpc;
        end
        else begin // first instr, may or may not be in stall ***** CHANGE LOGIC SO THAT CAN USE VALID/INVALID F2D FOR STALLING
            iCache.req = Valid(MemReq{addr: pc, op: Ld, data: ?});
        end

        f2d <= Valid(?);
        //////////////////////////
        // TODO: Your code here //
        //////////////////////////

    endrule

    method Word getPc = pc;
endmodule