import chisel3._
import chisel3.util._

class CSRFile(p: CoreParams = CoreParams()) extends Module {
  val io = IO(new Bundle {
    val cmd = Input(new CsrCommand(p))
    val readAddr = Input(UInt(12.W))
    val rdata = Output(UInt(p.xlen.W))
    val trap = Input(Bool())
    val trapPc = Input(UInt(p.xlen.W))
    val trapCause = Input(UInt(p.xlen.W))
    val satp = Output(UInt(p.xlen.W))
    val mtvec = Output(UInt(p.xlen.W))
    val mepc = Output(UInt(p.xlen.W))
  })

  val cycle = RegInit(0.U(p.xlen.W))
  val instret = RegInit(0.U(p.xlen.W))
  val mstatus = RegInit(0.U(p.xlen.W))
  val misa = RegInit("h800000000014112d".U(p.xlen.W))
  val mtvecReg = RegInit(0.U(p.xlen.W))
  val mepcReg = RegInit(0.U(p.xlen.W))
  val mcause = RegInit(0.U(p.xlen.W))
  val mtval = RegInit(0.U(p.xlen.W))
  val medeleg = RegInit(0.U(p.xlen.W))
  val mideleg = RegInit(0.U(p.xlen.W))
  val mie = RegInit(0.U(p.xlen.W))
  val mip = RegInit(0.U(p.xlen.W))
  val sstatus = RegInit(0.U(p.xlen.W))
  val stvec = RegInit(0.U(p.xlen.W))
  val sepc = RegInit(0.U(p.xlen.W))
  val scause = RegInit(0.U(p.xlen.W))
  val stval = RegInit(0.U(p.xlen.W))
  val sscratch = RegInit(0.U(p.xlen.W))
  val satpReg = RegInit(0.U(p.xlen.W))

  cycle := cycle + 1.U
  io.satp := satpReg
  io.mtvec := mtvecReg
  io.mepc := mepcReg

  val r = WireDefault(0.U(p.xlen.W))
  switch(io.readAddr) {
    is("h300".U) { r := mstatus }
    is("h301".U) { r := misa }
    is("h302".U) { r := medeleg }
    is("h303".U) { r := mideleg }
    is("h304".U) { r := mie }
    is("h305".U) { r := mtvecReg }
    is("h341".U) { r := mepcReg }
    is("h342".U) { r := mcause }
    is("h343".U) { r := mtval }
    is("h344".U) { r := mip }
    is("h100".U) { r := sstatus }
    is("h105".U) { r := stvec }
    is("h140".U) { r := sscratch }
    is("h141".U) { r := sepc }
    is("h142".U) { r := scause }
    is("h143".U) { r := stval }
    is("h180".U) { r := satpReg }
    is("hf14".U) { r := 0.U }
    is("hc00".U) { r := cycle }
    is("hc01".U) { r := 0.U }
    is("hc02".U) { r := instret }
  }
  io.rdata := r

  val writeOld = WireDefault(0.U(p.xlen.W))
  switch(io.cmd.addr) {
    is("h300".U) { writeOld := mstatus }
    is("h301".U) { writeOld := misa }
    is("h302".U) { writeOld := medeleg }
    is("h303".U) { writeOld := mideleg }
    is("h304".U) { writeOld := mie }
    is("h305".U) { writeOld := mtvecReg }
    is("h341".U) { writeOld := mepcReg }
    is("h342".U) { writeOld := mcause }
    is("h343".U) { writeOld := mtval }
    is("h344".U) { writeOld := mip }
    is("h100".U) { writeOld := sstatus }
    is("h105".U) { writeOld := stvec }
    is("h140".U) { writeOld := sscratch }
    is("h141".U) { writeOld := sepc }
    is("h142".U) { writeOld := scause }
    is("h143".U) { writeOld := stval }
    is("h180".U) { writeOld := satpReg }
    is("hf14".U) { writeOld := 0.U }
    is("hc00".U) { writeOld := cycle }
    is("hc01".U) { writeOld := 0.U }
    is("hc02".U) { writeOld := instret }
  }

  val operand = Mux(io.cmd.cmd(2), io.cmd.zimm, io.cmd.rs1)
  val next = MuxLookup(io.cmd.cmd(1, 0), writeOld)(Seq(
    1.U -> operand,
    2.U -> (writeOld | operand),
    3.U -> (writeOld & ~operand)
  ))

  when(io.cmd.valid && io.cmd.cmd =/= 0.U) {
    switch(io.cmd.addr) {
      is("h300".U) { mstatus := next }
      is("h302".U) { medeleg := next }
      is("h303".U) { mideleg := next }
      is("h304".U) { mie := next }
      is("h305".U) { mtvecReg := next }
      is("h341".U) { mepcReg := next }
      is("h342".U) { mcause := next }
      is("h343".U) { mtval := next }
      is("h344".U) { mip := next }
      is("h100".U) { sstatus := next }
      is("h105".U) { stvec := next }
      is("h140".U) { sscratch := next }
      is("h141".U) { sepc := next }
      is("h142".U) { scause := next }
      is("h143".U) { stval := next }
      is("h180".U) { satpReg := next }
    }
  }

  when(io.trap) {
    mepcReg := io.trapPc
    mcause := io.trapCause
  }
}