import chisel3._
import chisel3.util._

class ITLB(p: CoreParams = CoreParams()) extends Module {
  private val vpnBits = 27
  private val ppnBits = 44
  private val idxBits = log2Ceil(p.itlbEntries)

  val io = IO(new Bundle {
    val req = Input(new TlbReq(p))
    val resp = Output(new TlbResp(p))
    val refill = Input(new TlbRefill(p))
    val missVpn = Output(UInt(vpnBits.W))
  })

  val valid = RegInit(VecInit(Seq.fill(p.itlbEntries)(false.B)))
  val vpn = Reg(Vec(p.itlbEntries, UInt(vpnBits.W)))
  val ppn = Reg(Vec(p.itlbEntries, UInt(ppnBits.W)))
  val level = Reg(Vec(p.itlbEntries, UInt(2.W)))
  val flags = Reg(Vec(p.itlbEntries, UInt(8.W)))
  val repl = RegInit(0.U(idxBits.W))

  val reqVpn = io.req.vaddr(38, 12)
  val pageOff = io.req.vaddr(11, 0)
  val hitVec = VecInit((0 until p.itlbEntries).map(i => valid(i) && vpn(i) === reqVpn))
  val hit = io.req.valid && hitVec.asUInt.orR
  val hitIdx = OHToUInt(hitVec)
  val x = flags(hitIdx)(3)
  val pageFault = hit && !x

  io.resp.hit := hit && !pageFault
  io.resp.miss := io.req.valid && !hit
  io.resp.paddr := Cat(ppn(hitIdx), pageOff)
  io.resp.pageFault := pageFault
  io.resp.accessFault := false.B
  io.missVpn := reqVpn

  when(io.refill.valid) {
    valid(repl) := true.B
    vpn(repl) := io.refill.vpn
    ppn(repl) := io.refill.ppn
    level(repl) := io.refill.level
    flags(repl) := io.refill.flags
    repl := repl + 1.U
  }
}