package BACK

import chisel3._
import chisel3.util._

import org.chipsalliance.cde.config._
import freechips.rocketchip.diplomacy._
import freechips.rocketchip.tilelink._


abstract class TLCDRBase()(implicit p: Parameters) extends BackModule {
  class TLCDRIO extends Bundle{
    val overCLK = Input(Bool())
    val dDatIn   = Input(Bool())
    val dDatOut  = Output(Bool())
    val uDatIn   = Input(Bool())
    val uDatOut  = Output(Bool())

    val mem = Flipped(new Mem_Port_Bundle)

    val interrupt = Output(Vec(4, Bool()))
  }

  val io: TLCDRIO = IO(new TLCDRIO)

  val dirSel = RegInit(false.B)

  val CDRIn = Module(new CDRIn)
  CDRIn.io.serDat  := Mux( ~dirSel, io.dDatIn, io.uDatIn )
  CDRIn.io.overCLK := io.overCLK

  val CDROut = Module(new CDROut)
  io.dDatOut := ~dirSel & CDROut.io.serDat
  io.uDatOut :=  dirSel & CDROut.io.serDat



  val txFifo = Module(new Queue(UInt(32.W), 4))
  val rxFifo = Module(new Queue(UInt(32.W), 4))

}


trait TLCDRTx{ this: TLCDRBase =>

  val isTLReadStatus        = io.mem.fire & io.mem.addr(4,0) === "h0".U & ( io.mem.wstrb === 0.U )
  val isTLWriteTxLen        = io.mem.fire & io.mem.addr(4,0) === "h4".U & ( io.mem.wstrb =/= 0.U )
  val isTLWriteTxFifo       = io.mem.fire & io.mem.addr(4,0) === "h8".U & ( io.mem.wstrb =/= 0.U )


  val txLen     = RegInit(0.U(12.W))
  val isTxBusy  = RegInit(false.B)
  val isTxFull  = ~txFifo.io.enq.ready


  val intTxEnd      = ShiftRegister( CDROut.io.axis.fire & CDROut.io.axis.bits.tlast, 5, false.B, true.B)
  val intTxFifoFull = ~RegNext(isTxFull, false.B)  & isTxFull
  val intTxFinish   = CDROut.io.axis.fire & CDROut.io.axis.bits.tlast

  val txCnt = Reg(UInt(2.W))

  val txData = Reg(UInt(32.W))

  when( txFifo.io.deq.fire ){ //txCnt === 3.U
    txData := txFifo.io.deq.bits
    when( isTxBusy ){
      assert(CDROut.io.axis.fire)      
    }

  } .elsewhen( CDROut.io.axis.fire ){
    txData := txData << 8
  }
  
  txFifo.io.deq.ready := 
    Mux( isTxBusy, CDROut.io.axis.fire & txCnt === 3.U, true.B )

  txFifo.reset := reset.asBool | isTLWriteTxLen

  CDROut.io.axis.bits.tdata := txData.head(8)

  // PriorityMux(Seq(
  //   (txCnt === 0.U) -> txData(31, 24),
  //   (txCnt === 1.U) -> txData(23, 16),
  //   (txCnt === 2.U) -> txData(15,  8),
  //   (txCnt === 3.U) -> txData( 7,  0),
  // ))

  CDROut.io.axis.bits.tlast := ~txFifo.io.deq.valid & txCnt === 3.U
  CDROut.io.axis.bits.tuser := false.B
  CDROut.io.axis.valid := isTxBusy


  when( isTLWriteTxLen ){
    isTxBusy := false.B
  }.elsewhen( txFifo.io.deq.fire & ~isTxBusy ){
    isTxBusy := true.B
  } .elsewhen( CDROut.io.axis.fire & CDROut.io.axis.bits.tlast ){
    isTxBusy := false.B
  }


  when( isTLWriteTxLen ){
    txCnt := 0.U
  } .elsewhen( CDROut.io.axis.fire ){
    txCnt := txCnt + 1.U
  }


  when( isTLWriteTxLen ){
    txLen := io.mem.wdata
  } .elsewhen( CDROut.io.axis.fire ){
    txLen := txLen - 1.U
  }

  txFifo.io.enq.bits  := io.mem.wdata
  txFifo.io.enq.valid := isTLWriteTxFifo & txLen =/= 0.U


}


trait TLCDRRx{ this: TLCDRBase => 

  val isTLWriteSoftReset = io.mem.fire & io.mem.addr(4,0) === "h10".U & ( (io.mem.wstrb =/= 0.U) )
  // val isTLReadRxStatus     = io.mem.fire & io.mem.addr(4,0) === "h14".U & ( io.mem.wstrb === 0.U )
  val isTLReadRxLen        = io.mem.fire & io.mem.addr(4,0) === "h18".U & ( io.mem.wstrb === 0.U )
  val isTLReadRxFifo       = io.mem.fire & io.mem.addr(4,0) === "h1c".U & ( io.mem.wstrb === 0.U )

  val rxLen     = RegInit(0.U(12.W))
  // val isRxBusy  = RegInit(false.B)
  val isRxValid = rxFifo.io.enq.fire
  val isRxError = RegInit(false.B)
  val isAxisEnd = RegInit(false.B)
  val isRxEnd   = isAxisEnd & ~rxFifo.io.deq.valid

  val intRxError  = ~RegNext(isRxError, false.B) & isRxError
  // val intRxValid  = ~RegNext(isRxValid, false.B) & isRxValid
  val intRxStart = rxLen === 0.U & CDRIn.io.axis.fire
  val intRxEnd   = ~RegNext(isRxEnd, false.B) & isRxEnd


  val rxCnt = RegInit(0.U(2.W))
  val rxData = Reg(UInt(24.W))

  when(isTLWriteSoftReset ){
    isAxisEnd := false.B
  } .elsewhen( CDRIn.io.axis.fire & CDRIn.io.axis.bits.tlast ){
    isAxisEnd := true.B
  }

  
  when( isTLWriteSoftReset ){
    rxCnt := 0.U
  } .elsewhen( CDRIn.io.axis.fire ){
    rxCnt :=  rxCnt + 1.U

    rxData := Cat( rxData, CDRIn.io.axis.bits.tdata )

    // PriorityMux(Seq(
    //   (rxCnt === 0.U) -> Cat( CDRIn.io.axis.bits.tdata, 0.U(16.W) /*,0.U(8.W)*/ ), 
    //   (rxCnt === 1.U) -> Cat( rxData(23,16), CDRIn.io.axis.bits.tdata, 0.U(8.W) /*,0.U(8.W)*/ ),
    //   (rxCnt === 2.U) -> Cat( rxData(23, 8), CDRIn.io.axis.bits.tdata /*,0.U(8.W)*/         ),
    //   // (rxCnt === 3.U) -> Cat( rxData(31,8) , CDRIn.io.axis.bits.tdata  ),
    // ))
  }

  when( isTLWriteSoftReset ){
    rxLen := 0.U
  } .elsewhen( CDRIn.io.axis.fire ){
    rxLen := rxLen + 1.U
  }

  rxFifo.io.enq.valid := CDRIn.io.axis.fire & rxCnt === 3.U
  rxFifo.io.enq.bits  := Cat( rxData(23, 0) , CDRIn.io.axis.bits.tdata  )
  
  // PriorityMux(Seq(
  //     (rxCnt === 0.U) -> Cat( CDRIn.io.axis.bits.tdata, 0.U(24.W) ),
  //     (rxCnt === 1.U) -> Cat( rxData(23,16), CDRIn.io.axis.bits.tdata, 0.U(16.W) ),
  //     (rxCnt === 2.U) -> Cat( rxData(23, 8), CDRIn.io.axis.bits.tdata, 0.U(8.W ) ),
  //     (rxCnt === 3.U) -> Cat( rxData(23, 0) , CDRIn.io.axis.bits.tdata  ),
  //   ))
  assert( ~(( CDRIn.io.axis.fire & CDRIn.io.axis.bits.tlast) & rxCnt =/= 3.U), "Assert Failed! Rx must 4-byte Align!" )

  CDRIn.io.axis.ready := true.B

  rxFifo.reset := reset.asBool | isTLWriteSoftReset

  rxFifo.io.deq.ready := isTLReadRxFifo


  when( isTLWriteSoftReset ){
    isRxError := false.B
  } .elsewhen( rxFifo.io.enq.valid & ~rxFifo.io.enq.ready ){
    isRxError := true.B
    printf("Warning, RxFifo Overflow!\n")
  }



}

class TLCDR()(implicit p: Parameters) extends TLCDRBase() with TLCDRTx with TLCDRRx{

  val isWriteTransDir       = io.mem.fire & io.mem.addr(4,0) === "hc".U & io.mem.wstrb =/= 0.U





  when( isWriteTransDir ){
    dirSel := io.mem.wdata
  }

  io.mem.ready := 
    ( io.mem.wstrb =/= 0.U & txFifo.io.enq.ready) | 
    ( io.mem.wstrb === 0.U & Mux( io.mem.addr(4,0) === "h1c".U, rxFifo.io.deq.valid, true.B ) )



  io.interrupt(0) := intTxEnd
  io.interrupt(1) := intRxError
  io.interrupt(2) := intRxStart
  io.interrupt(3) := intRxEnd


  io.mem.rdata :=
    Mux1H(Seq(
      isTLReadStatus -> Cat( isRxValid, isRxError, isTxBusy, isTxFull, false.B),
      isTLReadRxLen  -> rxLen,
      isTLReadRxFifo -> rxFifo.io.deq.bits
    ))




}