import chisel3._
import chisel3.util._
import _root_.circt.stage.ChiselStage

class Core(p: CoreParams = CoreParams()) extends Module {
  val io = IO(new Bundle {
    val imem_req_valid = Output(Bool())
    val imem_req_bits = Output(UInt(p.xlen.W))
    val imem_resp_valid = Input(Bool())
    val imem_resp_bits_0 = Input(UInt(32.W))
    val imem_resp_bits_1 = Input(UInt(32.W))

    val dmem_req_valid = Output(Bool())
    val dmem_req_bits_addr = Output(UInt(p.xlen.W))
    val dmem_req_bits_data = Output(UInt(p.xlen.W))
    val dmem_req_bits_isStore = Output(Bool())
    val dmem_req_bits_size = Output(UInt(3.W))
    val dmem_resp_valid = Input(Bool())
    val dmem_resp_bits = Input(UInt(p.xlen.W))
  })

  if (p.useOoOBackend) {
    val frontend = Module(new Frontend(p))
    val id = Module(new IDStage(p))
    val backend = Module(new OoOBackend(p))

    frontend.io.redirectValid := backend.io.flush
    frontend.io.redirectPc := backend.io.redirectPc
    frontend.io.invalidateICache := backend.io.invalidateICache
    frontend.io.imemRespValid := io.imem_resp_valid
    frontend.io.imemRespBits(0) := io.imem_resp_bits_0
    frontend.io.imemRespBits(1) := io.imem_resp_bits_1
    frontend.io.branchUpdate := 0.U.asTypeOf(new BranchUpdate(p))

    val fetchValid = RegInit(false.B)
    val fetchReg = Reg(new FetchPacket(p))
    val fetchReady = !fetchValid || backend.io.decodeReady
    frontend.io.outReady := fetchReady
    when(backend.io.flush) {
      fetchValid := false.B
    }.elsewhen(fetchReady) {
      fetchValid := frontend.io.outValid
      fetchReg := frontend.io.out
    }

    id.io.inValid := fetchValid
    id.io.in := fetchReg

    backend.io.decodeValid := id.io.outValid
    backend.io.decode := id.io.out
    backend.io.dmemRespValid := io.dmem_resp_valid
    backend.io.dmemRespData := io.dmem_resp_bits
    backend.io.satp := 0.U

    io.imem_req_valid := frontend.io.imemReqValid
    io.imem_req_bits := frontend.io.imemReqAddr
    io.dmem_req_valid := backend.io.dmemReqValid
    io.dmem_req_bits_addr := backend.io.dmemReq.addr
    io.dmem_req_bits_data := backend.io.dmemReq.data
    io.dmem_req_bits_isStore := backend.io.dmemReq.isStore
    io.dmem_req_bits_size := backend.io.dmemReq.size
  } else {

  val sFetch :: sExec :: sLoadWait :: Nil = Enum(3)
  val state = RegInit(sFetch)
  val pc = RegInit(Consts.ResetVector)
  val instReg = RegInit(0.U(32.W))
  val pcReg = RegInit(Consts.ResetVector)

  val regs = RegInit(VecInit(Seq.fill(32)(0.U(p.xlen.W))))

  val decoder = Module(new Decoder(p))
  decoder.io.pc := pcReg
  decoder.io.inst := instReg
  val dec = decoder.io.out

  val alu = Module(new ALU(p))
  val branch = Module(new BranchUnit(p))
  val csr = Module(new CSRFile(p))

  def regRead(addr: UInt): UInt = Mux(addr === 0.U, 0.U, regs(addr))
  def lowLoad(data: UInt, size: UInt, signed: Bool): UInt = {
    val b = data(7, 0)
    val h = data(15, 0)
    val w = data(31, 0)
    MuxLookup(size, data)(Seq(
      0.U -> Mux(signed, Consts.signExtend(b, 8), b),
      1.U -> Mux(signed, Consts.signExtend(h, 16), h),
      2.U -> Mux(signed, Consts.signExtend(w, 32), w),
      3.U -> data
    ))
  }

  val src1 = regRead(dec.rs1)
  val src2 = regRead(dec.rs2)
  val aluB = Mux(dec.isOpImm || dec.isJalr || dec.isLoad, dec.immI, src2)

  alu.io.fn := dec.aluFn
  alu.io.a := src1
  alu.io.b := aluB
  alu.io.isWord := dec.isWord
  branch.io.funct3 := dec.funct3
  branch.io.a := src1
  branch.io.b := src2

  csr.io.cmd.valid := state === sExec && dec.isSystem && dec.funct3 =/= 0.U &&
    !(dec.funct3(1) && dec.rs1 === 0.U)
  csr.io.cmd.addr := instReg(31, 20)
  csr.io.cmd.cmd := dec.funct3
  csr.io.cmd.rs1 := src1
  csr.io.cmd.zimm := dec.rs1
  csr.io.readAddr := instReg(31, 20)
  val isEcall = instReg === "h00000073".U
  val isEbreak = instReg === "h00100073".U
  val isMret = instReg === "h30200073".U
  val takeTrap = state === sExec && (isEcall || isEbreak)
  csr.io.trap := takeTrap
  csr.io.trapPc := pcReg
  csr.io.trapCause := Mux(isEbreak, 3.U, 11.U)

  val loadAddr = src1 + dec.immI
  val storeAddr = src1 + dec.immS
  val pendingLoadAddr = Reg(UInt(p.xlen.W))
  val pendingLoadRd = Reg(UInt(5.W))
  val pendingLoadSize = Reg(UInt(3.W))
  val pendingLoadSigned = Reg(Bool())

  val isAmo = instReg(6, 0) === "b0101111".U
  val amoOp = instReg(31, 27)
  val amoLoadLike = isAmo && amoOp === "b00010".U
  val amoStoreLike = isAmo && amoOp === "b00011".U

  val memReqInExec = state === sExec && (dec.isLoad || dec.isStore)
  io.imem_req_valid := state === sFetch
  io.imem_req_bits := pc
  io.dmem_req_valid := memReqInExec || state === sLoadWait
  io.dmem_req_bits_addr := Mux(state === sLoadWait, pendingLoadAddr, Mux(dec.isStore, storeAddr, loadAddr))
  io.dmem_req_bits_data := src2
  io.dmem_req_bits_isStore := state === sExec && Mux(isAmo, !amoLoadLike, dec.isStore)
  io.dmem_req_bits_size := Mux(state === sLoadWait, pendingLoadSize, dec.memWidth)

  val branchTarget = pcReg + dec.immB
  val jalTarget = pcReg + dec.immJ
  val jalrTarget = (src1 + dec.immI) & (~1.U(p.xlen.W))
  val branchTaken = dec.isBranch && branch.io.taken
  val nextPc = Mux(dec.isJal, jalTarget,
    Mux(dec.isJalr, jalrTarget,
      Mux(branchTaken, branchTarget, pcReg + 4.U)))

  val execWriteData = WireDefault(alu.io.out)
  when(dec.isLui) {
    execWriteData := dec.immU
  }.elsewhen(dec.isAuipc) {
    execWriteData := pcReg + dec.immU
  }.elsewhen(dec.isJal || dec.isJalr) {
    execWriteData := pcReg + 4.U
  }.elsewhen(dec.isSystem && dec.funct3 =/= 0.U) {
    execWriteData := csr.io.rdata
  }

  when(state === sFetch) {
    when(io.imem_resp_valid) {
      instReg := io.imem_resp_bits_0
      pcReg := pc
      state := sExec
    }
  }.elsewhen(state === sExec) {
    when(takeTrap) {
      pc := csr.io.mtvec
      state := sFetch
    }.elsewhen(isMret) {
      pc := csr.io.mepc
      state := sFetch
    }.elsewhen(dec.isLoad && !amoStoreLike) {
      when(io.dmem_resp_valid) {
        when(dec.rd =/= 0.U) {
          regs(dec.rd) := lowLoad(io.dmem_resp_bits, dec.memWidth, dec.memSigned)
        }
        pc := pcReg + 4.U
        state := sFetch
      }.otherwise {
        pendingLoadAddr := loadAddr
        pendingLoadRd := dec.rd
        pendingLoadSize := dec.memWidth
        pendingLoadSigned := dec.memSigned
        state := sLoadWait
      }
    }.elsewhen(dec.isStore || amoStoreLike) {
      pc := pcReg + 4.U
      state := sFetch
    }.otherwise {
      when(dec.writesRd && dec.rd =/= 0.U) {
        regs(dec.rd) := execWriteData
      }
      pc := nextPc
      state := sFetch
    }
  }.elsewhen(state === sLoadWait) {
    when(io.dmem_resp_valid) {
      when(pendingLoadRd =/= 0.U) {
        regs(pendingLoadRd) := lowLoad(io.dmem_resp_bits, pendingLoadSize, pendingLoadSigned)
      }
      pc := pcReg + 4.U
      state := sFetch
    }
  }

  }
}

object Core extends App {
  ChiselStage.emitSystemVerilogFile(
    new Core(),
    args = Array("--target-dir", "generated"),
    firtoolOpts = Array("-disable-all-randomization", "-strip-debug-info")
  )
}

object CoreOoO extends App {
  ChiselStage.emitSystemVerilogFile(
    new Core(CoreParams(useOoOBackend = true)),
    args = Array("--target-dir", "generated-ooo"),
    firtoolOpts = Array("-disable-all-randomization", "-strip-debug-info")
  )
}