import chisel3._
import chisel3.util._

class ICache(p: CoreParams = CoreParams()) extends Module {
  private val lineInsts = p.fetchWidth
  private val lineBytes = lineInsts * 4
  private val sets = p.iCacheBytes / (p.iCacheWays * lineBytes)
  private val setBits = log2Ceil(sets)
  private val instBits = log2Ceil(lineInsts)
  private val offsetBits = log2Ceil(lineBytes)

  val io = IO(new Bundle {
    val reqValid = Input(Bool())
    val reqAddr = Input(UInt(p.xlen.W))
    val reqPc = Input(UInt(p.xlen.W))
    val flush = Input(Bool())
    val invalidate = Input(Bool())
    val respReady = Input(Bool())
    val memReqValid = Output(Bool())
    val memReqAddr = Output(UInt(p.xlen.W))
    val memRespValid = Input(Bool())
    val memRespBits = Input(Vec(p.fetchWidth, UInt(32.W)))
    val respValid = Output(Bool())
    val resp = Output(new FetchPacket(p))
    val miss = Output(Bool())
  })

  def setIndex(addr: UInt): UInt = addr(offsetBits + setBits - 1, offsetBits)
  def instIndex(addr: UInt): UInt = addr(offsetBits - 1, 2)
  def tag(addr: UInt): UInt = addr(p.xlen - 1, offsetBits + setBits)
  def alignedFetchAddr(addr: UInt): UInt = Cat(addr(p.xlen - 1, offsetBits), 0.U(offsetBits.W))

  val valid = RegInit(VecInit(Seq.fill(sets)(VecInit(Seq.fill(p.iCacheWays)(
    VecInit(Seq.fill(lineInsts)(false.B)))))))
  val tags = SyncReadMem(sets, Vec(p.iCacheWays, UInt((p.xlen - offsetBits - setBits).W)))
  val data = SyncReadMem(sets, Vec(p.iCacheWays, Vec(lineInsts, UInt(32.W))))
  val repl = RegInit(VecInit(Seq.fill(sets)(0.U(log2Ceil(p.iCacheWays).W))))

  val sIdle :: sLookup :: sResp :: sMiss :: Nil = Enum(4)
  val state = RegInit(sIdle)

  val lookupAddr = Reg(UInt(p.xlen.W))
  val lookupPc = Reg(UInt(p.xlen.W))
  val lookupSet = Reg(UInt(setBits.W))
  val lookupInst = Reg(UInt(instBits.W))
  val lookupValidRow = Reg(Vec(p.iCacheWays, Vec(lineInsts, Bool())))

  val missAddr = Reg(UInt(p.xlen.W))
  val missPc = Reg(UInt(p.xlen.W))
  val missSet = Reg(UInt(setBits.W))
  val missInst = Reg(UInt(instBits.W))
  val missWay = Reg(UInt(log2Ceil(p.iCacheWays).W))
  val missRefillExisting = Reg(Bool())
  val missTagRow = Reg(Vec(p.iCacheWays, UInt((p.xlen - offsetBits - setBits).W)))
  val missDataRow = Reg(Vec(p.iCacheWays, Vec(lineInsts, UInt(32.W))))
  val missValidRow = Reg(Vec(p.iCacheWays, Vec(lineInsts, Bool())))
  val missReqSent = RegInit(false.B)

  val respReg = Reg(new FetchPacket(p))

  val reqSet = setIndex(io.reqAddr)
  val reqInst = instIndex(io.reqAddr)
  val readFire = state === sIdle && io.reqValid && !io.flush
  val readTags = tags.read(reqSet, readFire)
  val readData = data.read(reqSet, readFire)

  val tagHitVec = VecInit((0 until p.iCacheWays).map(w => lookupValidRow(w).asUInt.orR && readTags(w) === tag(lookupAddr)))
  val hitVec = VecInit((0 until p.iCacheWays).map(w => tagHitVec(w) && lookupValidRow(w)(lookupInst)))
  val hit = hitVec.asUInt.orR
  val hitWay = OHToUInt(hitVec)
  val tagHit = tagHitVec.asUInt.orR
  val tagHitWay = OHToUInt(tagHitVec)
  val lastInst = (lineInsts - 1).U(instBits.W)
  val nextInst = lookupInst + 1.U
  val lookupCanFetchPair = lookupInst =/= lastInst
  val lookupLane1Valid = lookupCanFetchPair && lookupValidRow(hitWay)(nextInst)

  val lookupResp = WireDefault(0.U.asTypeOf(new FetchPacket(p)))
  lookupResp.pc := lookupPc
  lookupResp.inst(0) := readData(hitWay)(lookupInst)
  lookupResp.inst(1) := Mux(lookupLane1Valid, readData(hitWay)(nextInst), 0.U)
  lookupResp.laneValid(0) := true.B
  lookupResp.laneValid(1) := lookupLane1Valid
  lookupResp.predictedTaken := false.B
  lookupResp.predictedTarget := lookupPc + (4 * p.fetchWidth).U

  val missCanFetchPair = missInst =/= lastInst
  val missResp = WireDefault(0.U.asTypeOf(new FetchPacket(p)))
  missResp.pc := missPc
  missResp.inst(0) := io.memRespBits(0)
  missResp.inst(1) := Mux(missCanFetchPair, io.memRespBits(1), 0.U)
  missResp.laneValid(0) := true.B
  missResp.laneValid(1) := missCanFetchPair
  missResp.predictedTaken := false.B
  missResp.predictedTarget := missPc + (4 * p.fetchWidth).U

  io.memReqValid := state === sMiss && !missReqSent
  io.memReqAddr := Mux(state === sMiss,
    Mux(missCanFetchPair, alignedFetchAddr(missAddr), missAddr),
    Mux(reqInst =/= lastInst, alignedFetchAddr(io.reqAddr), io.reqAddr))
  io.respValid := (state === sLookup && hit) || state === sResp || (state === sMiss && io.memRespValid)
  io.resp := Mux(state === sResp, respReg,
    Mux(state === sMiss && io.memRespValid, missResp, lookupResp))
  io.miss := state === sLookup && !hit || state === sMiss

  when(io.invalidate) {
    valid := VecInit(Seq.fill(sets)(VecInit(Seq.fill(p.iCacheWays)(
      VecInit(Seq.fill(lineInsts)(false.B))))))
    state := sIdle
    missReqSent := false.B
  }.elsewhen(io.flush) {
    state := sIdle
    missReqSent := false.B
  }.elsewhen(state === sIdle) {
    when(io.reqValid) {
      lookupAddr := io.reqAddr
      lookupPc := io.reqPc
      lookupSet := reqSet
      lookupInst := reqInst
      lookupValidRow := valid(reqSet)
      state := sLookup
    }
  }.elsewhen(state === sLookup) {
    when(hit) {
      repl(lookupSet) := hitWay
      when(io.respReady) {
        state := sIdle
      }.otherwise {
        respReg := lookupResp
        state := sResp
      }
    }.otherwise {
      missAddr := lookupAddr
      missPc := lookupPc
      missSet := lookupSet
      missInst := lookupInst
      missWay := Mux(tagHit, tagHitWay, repl(lookupSet))
      missRefillExisting := tagHit
      missTagRow := readTags
      missDataRow := readData
      missValidRow := lookupValidRow
      missReqSent := false.B
      state := sMiss
    }
  }.elsewhen(state === sResp) {
    when(io.respReady) {
      state := sIdle
    }
  }.elsewhen(state === sMiss) {
    when(!missReqSent) {
      missReqSent := true.B
    }
    when(io.memRespValid) {
      val tagWrite = Wire(Vec(p.iCacheWays, UInt((p.xlen - offsetBits - setBits).W)))
      val dataWrite = Wire(Vec(p.iCacheWays, Vec(lineInsts, UInt(32.W))))
      val validWrite = Wire(Vec(p.iCacheWays, Vec(lineInsts, Bool())))
      tagWrite := missTagRow
      dataWrite := missDataRow
      validWrite := missValidRow
      tagWrite(missWay) := tag(missAddr)
      when(!missRefillExisting) {
        for (i <- 0 until lineInsts) {
          validWrite(missWay)(i) := false.B
        }
      }
      dataWrite(missWay)(missInst) := io.memRespBits(0)
      validWrite(missWay)(missInst) := true.B
      when(missCanFetchPair) {
        dataWrite(missWay)(missInst + 1.U) := io.memRespBits(1)
        validWrite(missWay)(missInst + 1.U) := true.B
      }
      valid(missSet) := validWrite
      tags.write(missSet, tagWrite)
      data.write(missSet, dataWrite)
      when(!missRefillExisting) {
        repl(missSet) := missWay + 1.U
      }
      when(io.respReady) {
        state := sIdle
      }.otherwise {
        respReg := missResp
        state := sResp
      }
    }
  }
}