import CacheTypes;
import CacheHelpers;
import MainMemory;

// ReqStatus (defined in CacheTypes.ms) is used to keep track of the state of the current request
//typedef enum {
//    Ready,         // The cache is ready for a new request
//    Lookup,        // Issued a lookup to tag/data arrays
//    Writeback,     // In main memory access for dirty writeback
//    Fill           // In main memory access for requested data
//} ReqStatus;
//
// Possible flows:
//   HIT: Ready -> Lookup -> Ready
//   MISS, line is clean: Ready -> Lookup -> Fill
//   MISS, line is dirty: Ready -> Lookup -> Writeback -> Fill

// You may find the following type synonyms helpful to access the tag/data/status arrays
typedef SRAMReq#(logCacheSets, CacheTag) TagReq;
typedef SRAMReq#(logCacheSets, Line) DataReq;
typedef SRAMReq#(logCacheSets, CacheStatus) StatusReq;

// TODO: Complete the implementation of TwoWayCache
// NOTE: Implementing this module requires about 20 new lines of code from DirectMappedCache
module TwoWayCache(MainMemory mainMem);
    // SRAM arrays, one element per way (e.g., dataArrays[0] is the data array of way 0).
    Vector#(2, SRAM#(logCacheSets, Line)) dataArrays;
    Vector#(2, SRAM#(logCacheSets, CacheTag)) tagArrays;
    Vector#(2, SRAM#(logCacheSets, CacheStatus)) statusArrays;

    // LRU bits, one per set. We implement it as registers instead of an SRAM
    // array, because with small caches this is just a few bits (e.g., 64), and
    // SRAMs make sense only for larger memories. This also makes the code
    // simpler, and will make the cache easier to pipeline in the future.
    Vector#(cacheSets, RegU#(Bit#(1))) lru;

    // Registers for holding the current state of the cache and how far along
    // it is in processing a request.
    RegU#(MemReq) curReq;
    Reg#(ReqStatus) state(Ready);

    // Hit/miss counters
    Reg#(Word) hits(0);
    Reg#(Word) misses(0);

    input Maybe#(MemReq) req default = Invalid;

    //TODO return True if the cache can accept a new request
    method Bool reqEnabled = (state == Ready) ? True : False;

    rule tick;
        if (state == Ready && isValid(req)) begin
            //TODO Your code here
            let request = fromMaybe(?, req);
            tagArrays[0].req = Valid(TagReq {write: False, addr: getIndex(request.addr), data:?});
            tagArrays[1].req = Valid(TagReq {write: False, addr: getIndex(request.addr), data:?});
            statusArrays[0].req = Valid(StatusReq {write: False, addr: getIndex(request.addr), data:?});
            statusArrays[1].req = Valid(StatusReq {write: False, addr: getIndex(request.addr), data:?});
            dataArrays[0].req = Valid(DataReq {write: False, addr: getIndex(request.addr), data:?});
            dataArrays[1].req = Valid(DataReq {write: False, addr: getIndex(request.addr), data:?});
            state <= Lookup;
            curReq <= request;

        end else if (state == Lookup) begin
            //TODO Your code here
            let status0 = fromMaybe(?, statusArrays[0].data);
            let status1 = fromMaybe(?, statusArrays[1].data);
            let tag0 = fromMaybe(?, tagArrays[0].data);
            let tag1 = fromMaybe(?, tagArrays[1].data);
            let data0 = fromMaybe(?, dataArrays[0].data);
            let data1 = fromMaybe(?, dataArrays[1].data);

            //HIT on Way 1
            if (tag0 == getTag(curReq.addr) && status0 != Invalid) begin
                if (curReq.op == St) begin
                    data0[getWordOffset(curReq.addr)] = curReq.data;
                    statusArrays[0].req = Valid(StatusReq {write: True, addr: getIndex(curReq.addr), data: Dirty});
                    dataArrays[0].req = Valid(DataReq {write: True, addr: getIndex(curReq.addr), data:data0});
                end
                state <= Ready;
                hits <= hits + 1;
                lru[getIndex(curReq.addr)] <= 1;
            end
            //HIT on Way 2
            else if (tag1 == getTag(curReq.addr) && status1 != Invalid) begin
                if (curReq.op == St) begin
                    data1[getWordOffset(curReq.addr)] = curReq.data;
                    statusArrays[1].req = Valid(StatusReq {write: True, addr: getIndex(curReq.addr), data: Dirty});
                    dataArrays[1].req = Valid(DataReq {write: True, addr: getIndex(curReq.addr), data:data1});
                end
                state <= Ready;
                hits <= hits + 1;
                lru[getIndex(curReq.addr)] <= 0;
            end

            //MISS
            else begin
                misses <= misses +1;
                let i = lru[getIndex(curReq.addr)];
                //Clean
                let status = fromMaybe(?, statusArrays[i].data);
                if (status == Clean || status == Invalid) begin
                    mainMem.req = Valid(LineReq {op: Ld, lineAddr: {getTag(curReq.addr), getIndex(curReq.addr)}, data:?});
                    state <= Fill;
                end
                //Dirty
                else begin
                    let tag = fromMaybe(?, tagArrays[i].data);
                    mainMem.req = Valid(LineReq {op: St, lineAddr: {tag, getIndex(curReq.addr)}, data:fromMaybe(?, dataArrays[i].data)});
                    state <= Writeback;
                end
                //lru[getIndex(curReq.addr)] <= ~i;
            end


        end else if (state == Writeback && mainMem.reqEnabled) begin
            //TODO Your code here
            LineReq memReq = LineReq {op: Ld, lineAddr: getLineAddr(curReq.addr), data:?};
            mainMem.req = Valid(memReq);
            state <= Fill;

        end else if (state == Fill && isValid(mainMem.data)) begin
            //TODO Your code here
            let data = fromMaybe(?, mainMem.data);
            let i = lru[getIndex(curReq.addr)];
            StatusReq statusReq = StatusReq {write: True, addr: getIndex(curReq.addr), data: Clean};
            if (curReq.op == St) begin
                data[getWordOffset(curReq.addr)] = curReq.data;
                statusReq.data = Dirty;
            end
            statusArrays[i].req = Valid(statusReq);
            tagArrays[i].req = Valid(TagReq {write: True, addr: getIndex(curReq.addr), data: getTag(curReq.addr)});
            dataArrays[i].req = Valid(DataReq {write: True, addr: getIndex(curReq.addr), data: data});
            state <= Ready;
            lru[getIndex(curReq.addr)] <= ~i;
        end

    endrule

    method Maybe#(Word) data;
        // This method should return a Valid output in only two cases:
        // 1. On a load hit (we're in the Lookup state, there's a hit, and
        //    curReq.op == Ld).
        // 2. On a fill for a load request (we're in the Fill state,
        //    mainMem.data is valid, and curReq.op == Ld).
        // In all other cases, the output should be Invalid
        //

        //TODO Your code here
        let status0 = fromMaybe(?, statusArrays[0].data);
        let status1 = fromMaybe(?, statusArrays[1].data);
        let tag0 = fromMaybe(?, tagArrays[0].data);
        let tag1 = fromMaybe(?, tagArrays[1].data);

        if (state == Fill && curReq.op == Ld && isValid(mainMem.data)) begin
            let line = fromMaybe(?, mainMem.data);
            return Valid(line[getWordOffset(curReq.addr)]);
        end
        else if (state == Lookup && curReq.op == Ld) begin
            if (status0 != Invalid && tag0 == getTag(curReq.addr)) begin
                let line = fromMaybe(?, dataArrays[0].data);
                return Valid(line[getWordOffset(curReq.addr)]);
            end
            else if (status1 != Invalid && tag1 == getTag(curReq.addr)) begin
                let line = fromMaybe(?, dataArrays[1].data);
                return Valid(line[getWordOffset(curReq.addr)]);
            end
            else return Invalid;
        end
	else return Invalid;
    endmethod
    method Bit#(32) getHits = hits;
    method Bit#(32) getMisses = misses;
endmodule