/***************************************************************************************
 * Copyright (c) 2020-2025 Institute of Computing Technology, Chinese Academy of Sciences
 * Copyright (c) 2025 Beijing Institute of Open Source Chip
 *
 * DiffTest is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 *
 * See the Mulan PSL v2 for more details.
 ***************************************************************************************/

package difftest

import chisel3._
import chisel3.experimental.prefix
import chisel3.reflect.DataMirror
import difftest.common.DifftestWiring
import difftest.fpga.HostEndpoint

class DifftestTopIO extends Bundle {
  val exit = Output(UInt(64.W))
  val step = Output(UInt(64.W))
  val perfCtrl = new PerfCtrlIO
  val logCtrl = new LogCtrlIO
  val uart = new UARTIO
}

class PerfCtrlIO extends Bundle {
  val clean = Input(Bool())
  val dump = Input(Bool())
}

class LogCtrlIO extends Bundle {
  val begin = Input(UInt(64.W))
  val end = Input(UInt(64.W))
  val level = Input(UInt(64.W))

  def enable(timer: UInt): Bool = {
    val en = WireInit(false.B)
    en := timer >= begin && timer < end
    en
  }
}

// UART IO: input/output through the simulation framework
class UARTIO extends Bundle {
  val out = new Bundle {
    val valid = Output(Bool())
    val ch = Output(UInt(8.W))
  }
  val in = new Bundle {
    val valid = Output(Bool())
    val ch = Input(UInt(8.W))
  }
}

trait HasDiffTestInterfaces extends ImplicitClock with ImplicitReset { this: RawModule =>
  def cpuName: Option[String] = None

  private[difftest] def dutClock = implicitClock
  private[difftest] def dutReset = implicitReset
  private[difftest] def dutIOs: Seq[(String, Data)] = {
    def portFilter(port: (String, Data)): Boolean = {
      port._1.contains("bore") || // Internal port created by BoringUtils, connected to difftest only
      (port._2 match {
        case u: UARTIO => true // UART port for simulation framework, connected to difftest only
        case b: Bundle => b.elements.toSeq.exists(portFilter)
        case other     => Seq(dutClock, dutReset).contains(other) // Already generated by Chisel module
      })
    }
    def portCollect(ports: Seq[(String, Data)]): Seq[(String, Data)] = {
      ports.filterNot(portFilter) ++
        ports.filter(portFilter).flatMap { case (root, data) =>
          data match {
            case u: UARTIO => Seq.empty
            case b: Bundle => portCollect(b.elements.toSeq.map { case (n, d) => (s"${root}_${n}", d) })
            case _         => Seq.empty
          }
        }
    }
    portCollect(DataMirror.modulePorts(this))
  }

  def connectTopIOs(difftest: DifftestTopIO): Unit = {}
}

// Top-level module for DiffTest simulation. Will be created by DifftestModule.top
class SimTop[T <: RawModule with HasDiffTestInterfaces](cpuGen: => T, modPrefix: Option[String]) extends Module {
  override def desiredName: String = modPrefix.getOrElse("") + super.desiredName
  val cpu = Module(cpuGen)
  cpu.dutClock := clock
  cpu.dutReset := reset.asTypeOf(cpu.dutReset)

  val cpuName = cpu.cpuName.getOrElse(cpu.getClass.getName.split("\\.").last)
  val gateway = DifftestModule.collect(cpuName)

  // IO: difftest_*
  val difftest = IO(new DifftestTopIO)

  difftest.exit := gateway.exit.getOrElse(0.U)
  difftest.step := gateway.step.getOrElse(0.U)
  difftest.uart := DontCare

  prefix("difftest") {
    // Required signals for LogPerfControl
    val timer = RegInit(0.U(64.W))
    timer := timer + 1.U
    dontTouch(timer)

    val log_enable = difftest.logCtrl.enable(timer)
    dontTouch(log_enable)

    val ref_clock = Option.when(gateway.refClock.isDefined)(IO(Input(Clock())))
    gateway.refClock.foreach(_ := ref_clock.get)

    // IO: difftest_fpga_*
    gateway.fpgaIO.foreach { fpgaIO =>
      val host = withClock(ref_clock) { Module(new HostEndpoint(fpgaIO.bits.getWidth)) }
      host.io.difftest <> fpgaIO
      val pcie_clock = IO(Input(Clock()))
      host.io.pcie_clock := pcie_clock

      val toHost = host.io.to_host_axis
      val to_host_axis = IO(chiselTypeOf(toHost))
      to_host_axis <> toHost
    }

    gateway.clockEnable.foreach { clockEnable =>
      val clock_enable = IO(Output(Bool()))
      clock_enable := clockEnable
    }
  }

  cpu.connectTopIOs(difftest)
  cpu.dutIOs.foreach { case (name, gen) =>
    val io = IO(chiselTypeOf(gen)).suggestName(name)
    dontTouch(gen)
    io <> gen
  }

  // There should not be anymore IOs
  require(DifftestWiring.isEmpty, s"pending wires left: ${DifftestWiring.getPending}")
}
