package vexiiriscv.misc

import spinal.core._
import spinal.lib.fsm.{State, StateMachine}
import spinal.lib._
import spinal.lib.misc.plugin.FiberPlugin
import vexiiriscv.Global
import vexiiriscv.execute.{CsrAccessPlugin, CsrListFilter, CsrRamAllocation, CsrRamPlugin, CsrRamService}
import vexiiriscv.riscv.{CSR, PrivilegeMode, Riscv}

import scala.collection.mutable.ArrayBuffer

/**
 * This plugin implement the performance counters in a very tricky way to save area. The RISC-V spec specified 64 bits per counter, and that is kinda expensive to implement FPGA.
 * So in this implementation, Only 8 bits registers are used for each counters, which are flushed into their 64 bits CSR ram when their MSB is set
 *
 * To figure out the actual performance events implement in VexiiRIscv, look at vexiiriscv.misc.PerformanceCounterService
 */
class PerformanceCounterPlugin(var additionalCounterCount : Int,
                               var bufferWidth : Int = 8,
                               var withSmcntrpmf : Boolean = true,
                               var withScountovf : Boolean = true,
                               var withShlcofideleg : Boolean = false) extends FiberPlugin with PerformanceCounterService{
  def counterCount = 2 + additionalCounterCount

  case class Spec(id : Int, event : Bool)
  val specs = ArrayBuffer[Spec]()
  override def createEventPort(id: Int) = specs.addRet(Spec(id, Bool())).event

  val logic = during setup new Area{
    val csr = host[CsrAccessPlugin]
    val ram = host[CsrRamPlugin]
    val priv = host[PrivilegedPlugin]
    val tp = host[TrapPlugin]
    val csrRetainer = csr.csrLock()
    val ramCsrRetainer = ram.csrLock()
    val ramPortRetainer = ram.portLock()
//    val trapLock = tp.trapLock()
    awaitBuild()


    assert(Global.HART_COUNT.get == 1)
    val withHigh = Riscv.XLEN.get == 32

    assert(Global.HART_COUNT.get == 1)
    val commitMask = Cat(host.list[CommitService].map(_.getCommitMask(0)))
    val ignoreNextCommit = RegInit(False) clearWhen (commitMask.orR)
    val commitCount = CountOne(commitMask) - U(ignoreNextCommit && commitMask.orR)

    val eventCycles = createEventPort(PerformanceCounterService.CYCLES)
    val eventInstructions = Vec.fill(widthOf(commitMask))(createEventPort(PerformanceCounterService.INSTRUCTIONS))

    var counterIdPtr = 0
    val privValue = priv.getPrivilege(0)
    val rawPrivValue = PrivilegeMode.mode(privValue)

    class Counter(withOverflow: Boolean = true) extends Area{
      val alloc = ram.ramAllocate(if (withHigh) 2 else 1)
      val value = Reg(UInt(bufferWidth bits)).init(0)
      val needFlush = value.msb
      val counterId = counterIdPtr
      counterIdPtr += 1

      val mcounteren = csr.readWrite(RegInit(False), CSR.MCOUNTEREN, counterId)
      val scounteren = priv.p.withSupervisor generate csr.readWrite(RegInit(False), CSR.SCOUNTEREN, counterId)
      val hcounteren = priv.p.withHypervisor generate csr.readWrite(RegInit(False), CSR.HCOUNTEREN, counterId)
      val mcountinhibit = csr.readWrite(RegInit(False), CSR.MCOUNTINHIBIT, counterId)

      val MINH    = RegInit(False)
      val SINH    = RegInit(False)
      val UINH    = RegInit(False)
      val VSINH   = RegInit(False)
      val VUINH   = RegInit(False)
      val OF      = withOverflow generate RegInit(False)
      val inhibit = CombInit(mcountinhibit)
      val ofRead  = withOverflow generate CombInit(OF)

      def mappingCfgCsr(eb: Int, eo: Int) = new Area {
        if (withOverflow) ofRead clearWhen(!mcounteren && rawPrivValue =/= PrivilegeMode.M)

        if (withOverflow) csr.readWrite(eb, 63-eo -> OF)
        csr.readWrite(eb, 62-eo -> MINH)
        inhibit.setWhen(privValue === PrivilegeMode.M && MINH)
        if (priv.p.withSupervisor) {
          csr.readWrite(eb, 61 - eo -> SINH)
          inhibit.setWhen(privValue === PrivilegeMode.S && SINH)
        }
        if (priv.p.withUser) {
          csr.readWrite(eb, 60 - eo -> UINH)
          inhibit.setWhen(privValue === PrivilegeMode.U && UINH)
        }
        if (priv.p.withHypervisor) {
          csr.readWrite(eb, 59 - eo -> VSINH, 58 - eo -> VUINH)
          inhibit.setWhen(privValue === PrivilegeMode.VS && VSINH)
          inhibit.setWhen(privValue === PrivilegeMode.VU && VUINH)
          if (withOverflow) ofRead clearWhen(!hcounteren && PrivilegeMode.isGuest(privValue))
        }
      }

    }
    case class Mapping(csrId : Int, alloc : CsrRamAllocation, offset : Int)
    val mappings = ArrayBuffer[Mapping]()
    val counters = new Area{
      val cycle = new Counter(withOverflow = false)
      counterIdPtr += 1 //skip time
      val instret = new Counter(withOverflow = false)
      val additionals = List.fill(additionalCounterCount)(new Counter)

      eventCycles := True
      eventInstructions := (commitMask.andMask(!ignoreNextCommit)).asBools

      cycle.value := cycle.value + (!cycle.mcountinhibit).asUInt
      instret.value := instret.value + RegNext(commitCount.andMask(!instret.mcountinhibit)).init(0)

      if (withSmcntrpmf) Riscv.XLEN.get match {
        case 32 => {
          cycle.mappingCfgCsr(CSR.MCYCLECFGH, 32)
          instret.mappingCfgCsr(CSR.MINSTRETCFGH, 32)
          csr.allowCsr(CsrListFilter(List(CSR.MCYCLECFG, CSR.MINSTRETCFG)))
        }
        case 64 => {
          cycle.mappingCfgCsr(CSR.MCYCLECFG, 0)
          instret.mappingCfgCsr(CSR.MINSTRETCFG, 0)
        }
      }

      mappings += Mapping(CSR.MCYCLE, cycle.alloc, 0)
      mappings += Mapping(CSR.UCYCLE, cycle.alloc, 0)
      mappings += Mapping(CSR.MINSTRET, instret.alloc, 0)
      mappings += Mapping(CSR.UINSTRET, instret.alloc, 0)
      for(i <- 0 until additionalCounterCount) {
        mappings += Mapping(CSR.MHPMCOUNTER3 + i, additionals(i).alloc, 0)
        mappings += Mapping(CSR.UHPMCOUNTER3 + i, additionals(i).alloc, 0)
      }
      if (withHigh) for(b <- mappings.toArray) mappings += Mapping(b.csrId+0x80, b.alloc, b.offset+1)

      for (m <- mappings) {
        ram.readWriteRam(m.csrId, m.alloc, m.offset)
      }
      val list = List(cycle, instret) ++ additionals
    }

    val dummyCsrs = for(hpmId <- 3+additionalCounterCount to 31;
                        csrId <- List(CSR.MHPMCOUNTER0, CSR.UHPMCOUNTER0, CSR.MHPMEVENT0) ++ (if(withHigh)List(CSR.MHPMCOUNTER0H, CSR.UHPMCOUNTER0H, CSR.MHPMEVENT0H) else Nil)) yield csrId + hpmId
    csr.allowCsr(CsrListFilter(dummyCsrs)) //As read zero


    ramCsrRetainer.release()
    val readPort = ram.ramReadPort(CsrRamService.priority.COUNTER)
    val writePort = ram.ramWritePort(CsrRamService.priority.COUNTER)
    ramPortRetainer.release()

    elaborationLock.await()

    val csrFilter = CsrListFilter(mappings.map(m => m.csrId))

    val interrupt = new Area {
      val ip, ie = RegInit(False)
      csr.readWrite(CSR.MIP, 13 -> ip)
      csr.readWrite(CSR.MIE, 13 -> ie)
      val sup = priv.implementSupervisor generate new Area {
        val deleg = RegInit(False)
        csr.readWrite(CSR.MIDELEG, 13 -> deleg)
        csr.read(CSR.SIP, 13 -> (ip & deleg))
        csr.writeWhen(ip, deleg, CSR.SIP, 13)
        csr.read(CSR.SIE, 13 -> (ie & deleg))
        csr.writeWhen(ie, deleg, CSR.SIE, 13)
      }

      val vs = (priv.implementHypervisor && withShlcofideleg) generate new Area {
        val deleg = RegInit(False)
        val active = sup.deleg && deleg

        if(withShlcofideleg && withScountovf) {
          deleg clearWhen(!sup.deleg)
          csr.read(sup.deleg && deleg, CSR.HIDELEG, 13)
          csr.writeWhen(deleg, sup.deleg, CSR.HIDELEG, 13)
        }

        csr.read(ip && active, CSR.VSIP, 13)
        csr.writeWhen(ip, active, CSR.VSIP, 13)
        csr.read(ie && active, CSR.VSIE, 13)
        csr.writeWhen(ie, active, CSR.VSIE, 13)

        priv.logic.harts(0).spec.addInterrupt(
          ip && ie && active,
          id = 13,
          privilege = PrivilegeMode.VS,
          delegators = List(
            Delegator(True, PrivilegeMode.M),
            Delegator(True, PrivilegeMode.S)
          )
        )
      }
      val nonGuestMode = (priv.implementHypervisor && withShlcofideleg).mux(!vs.active, True)

      priv.logic.harts(0).spec.addInterrupt(
        ip && ie && nonGuestMode,
        id = 13,
        privilege = priv.implementSupervisor.mux(PrivilegeMode.S, PrivilegeMode.M),
        delegators = priv.implementSupervisor.mux(List(Delegator(sup.deleg, PrivilegeMode.M)), Nil)
      )
    }

    val events = new Area {
      val selWidth = log2Up((specs.map(_.id) :+ 0).max + 1)
      val grouped = specs.groupByLinked(_.id)
      val sums = grouped.map{ case (id, specs) => id -> CountOne(RegNext(specs.map(_.event).asBits)) }
      val widthMax = sums.map(_._2.getWidth).max
    }

    val hpm = for(id <- 0 until 3+additionalCounterCount) yield (id >= 3) generate new Area{
      val counter = counters.additionals(id-3)
      val eventId = Reg(UInt(events.selWidth bits)) init(0)
      val overflowEvent = False
      counter.OF setWhen(overflowEvent)

      interrupt.ip.setWhen(overflowEvent && !counter.OF)

      val incr    = if(events.sums.isEmpty) U(0) else events.sums.map(e => e._2.andMask(eventId === e._1).resize(events.widthMax)).toList.reduceBalancedTree(_ | _)
      when(!counter.inhibit) {
        counter.value := counter.value + incr
      }
      csr.readWrite(CSR.MHPMEVENT0 + id, 0 -> eventId)
      if(withHigh) csr.allowCsr(CSR.MHPMEVENT0H + id)
      val (eb, eo) = Riscv.XLEN.get match {
        case 32 => (CSR.MHPMEVENT0H + id, 32)
        case 64 => (CSR.MHPMEVENT0  + id, 0)
      }
      counter.mappingCfgCsr(eb, eo)

      if(withScountovf && priv.implementSupervisor) csr.read(CSR.SCOUNTOVF, id -> counter.ofRead)
    }


    val fsm = new StateMachine{
      val IDLE, READ_LOW, CALC_LOW = new State
      val READ_HIGH, CALC_HIGH = withHigh generate new State
      val CSR_WRITE = new State
      setEntry(IDLE)


      val flusherCmd = Stream(new Bundle {
        val oh = Bits(counterCount bits)
      })
      val csrWriteCmd, csrReadCmd = Stream(new Bundle {
        val address = UInt(log2Up(counterCount+1) bits)
      })
      List(csrReadCmd, flusherCmd, csrWriteCmd).foreach(_.ready := False)


      ram.awaitMapping()
      val cmd = new Area {
        val flusher = Reg(Bool())
        val oh = Reg(Bits(counterCount bits))
        val address = counters.list.reader(oh)(c => U(c.alloc.at >> withHigh.toInt, ram.portAddressWidth - withHigh.toInt bits))
      }

      val done = isActive(IDLE)
      val ramReaded = Reg(UInt(Riscv.XLEN bits))
      val carry = Reg(Bool())
      val counterReaded = Reg(UInt(bufferWidth bits))

      readPort.valid := False
      readPort.address := (if(withHigh) cmd.address @@ carry else cmd.address)

      writePort.valid := False
      writePort.address := readPort.address
      writePort.data.assignDontCare()


      val calc = new Area{
        def low = ramReaded(ramReaded.high downto bufferWidth - 1) @@ counterReaded.resize(bufferWidth - 1)
        def high = ramReaded
        def blow = U(counterReaded.msb) << bufferWidth-1

        val a = if(withHigh) carry.mux(high, low) else low
        val b = if(withHigh) carry.mux(U(1), blow) else blow
        val sum = a +^ b
      }

      def doOverflow(): Unit = {
        hpm.drop(3).onMask(cmd.oh.drop(2)) { c =>
          c.overflowEvent := True
        }
      }

      val idleCsrAddress = csrReadCmd.valid.mux(csrReadCmd.address, csrWriteCmd.address)
      val holdCsrWrite = True
      when(holdCsrWrite){
        ram.holdCsrWrite()
      }
      IDLE whenIsActive{
        holdCsrWrite := False
        cmd.oh := B(for (c <- counters.list) yield idleCsrAddress === c.counterId)
        when(csrWriteCmd.valid) {
          goto(CSR_WRITE)
        }elsewhen(flusherCmd.valid){
          cmd.flusher := True
          cmd.oh := flusherCmd.oh
          flusherCmd.ready := True
          goto(READ_LOW)
        } elsewhen(csrReadCmd.valid){
          cmd.flusher := False
          csrReadCmd.ready := True
          goto(READ_LOW)
        }
        carry := False
      }

      CSR_WRITE whenIsActive {
        holdCsrWrite := False
        when(csr.bus.write.address(7) === False) {
          counters.list.onMask(cmd.oh) { c =>
            c.value := csr.bus.write.bits.asUInt.resized
            c.value.msb := False
          }
        }
        ignoreNextCommit setWhen (cmd.oh(1)) // && (if(withHigh) !csrWriteCmd.high else True)
        csrWriteCmd.ready := True
        goto(IDLE)
      }


      READ_LOW.whenIsActive{
        readPort.valid := True
        ramReaded := U(readPort.data)
        counterReaded := counters.list.reader(cmd.oh)(_.value)
        when(readPort.ready) {
          goto(CALC_LOW)
        }
      }

      CALC_LOW.whenIsActive {
        when(counterReaded.msb) {
          counters.list.onMask(cmd.oh)(_.value.msb := False)
        }
        writePort.valid := True
        writePort.data := B(calc.sum).resized
        when(writePort.ready) {
          goto(IDLE)
          if (withHigh) when(calc.sum.msb){
            carry := True
            goto(READ_HIGH)
          } else {
            when(calc.sum.msb){
              doOverflow()
            }
          }
        }
      }


      if(withHigh) READ_HIGH.whenIsActive{
        readPort.valid := True
        ramReaded := U(readPort.data)
        when(readPort.ready) {
          goto(CALC_HIGH)
        }

        CALC_HIGH.whenIsActive {
          writePort.valid := True
          writePort.data := B(calc.sum).resized
          when(writePort.ready) {
            when(calc.sum.msb) {
              doOverflow()
            }
            goto(IDLE)
          }
        }
      }

    }

    val flusher = new Area{
      val hits = B(counters.list.map(_.needFlush))
      val hit = hits.orR
      val oh = OHMasking.first(hits)

      fsm.flusherCmd.valid := hit
      fsm.flusherCmd.oh := oh
    }

    // Trap eventual bad accesses
    val csrDecode = new Area{
      val addr = csr.bus.decode.address(0, log2Up(counterCount + 1) bits)
      val mok = addr.muxListDc(counters.list.map(e => e.counterId -> e.mcounteren))
      val sok = priv.p.withSupervisor.mux(addr.muxListDc(counters.list.map(e => e.counterId -> e.scounteren)), True)
      val vok = priv.p.withHypervisor.mux(addr.muxListDc(counters.list.map(e => e.counterId -> e.hcounteren)), True)
      val privilege = priv.getPrivilege(csr.bus.decode.hartId)
      val privOk = ((privilege.asBits ^ B(1 << 2)) | (vok ## mok ## sok)).andR
      val hyperOK = PrivilegeMode.isGuest(privilege) && (privilege(1 downto 0).asBits | (mok ## (vok || sok))).andR
      csr.onDecode(csrFilter){ //TODO test
        when(csr.bus.decode.address(9 downto 8) === PrivilegeMode.U){
          when(csr.bus.decode.write){
            csr.bus.decode.doHostDenied()
          }
          when(!privOk){
            csr.bus.decode.doException()
          }
        }
        when(csr.bus.decode.address(9 downto 8) < PrivilegeMode.M && hyperOK){
          csr.bus.decode.doVirtual()
        }
      }
    }

    val csrRead = new Area {
      val fired = RegInit(False) setWhen(fsm.csrReadCmd.fire)
      val requested = csr.bus.read.valid && csr.readingCsr(csrFilter)
      fsm.csrReadCmd.valid := requested && !fired
      fsm.csrReadCmd.address := csr.bus.read.address(0, log2Up(counterCount+1) bits)

      when(requested){
        when(!fired || !fsm.done){
          ram.holdCsrRead()
        }
      }
      when(csr.bus.read.moving){
        fired := False
      }
    }

    val csrWrite = new Area{
      val fired = RegInit(False) setWhen(fsm.csrWriteCmd.fire)
      fsm.csrWriteCmd.valid := False
      fsm.csrWriteCmd.address := csr.bus.write.address(0, log2Up(counterCount+1) bits)

      csr.onWrite(csrFilter, false){
        when(!fired){
          fsm.csrWriteCmd.valid := True
          when(!fsm.csrWriteCmd.ready){
            csr.bus.write.doHalt()
          }
        }
      }
      when(csr.bus.write.moving){
        fired := False
      }
    }
    csrRetainer.release()
//    trapLock.release()

    if(priv.p.withDebug) when(priv.logic.harts(0).debugMode){
      commitCount := 0
    }
  }
}
