package chipyard.harness

import chisel3._
import chisel3.util._
import chisel3.experimental.DoubleParam
import scala.collection.mutable.{ArrayBuffer, LinkedHashMap}
import freechips.rocketchip.diplomacy.{LazyModule}
import org.chipsalliance.cde.config.{Field, Parameters, Config}
import freechips.rocketchip.util.{ResetCatchAndSync}
import freechips.rocketchip.prci._

import chipyard.harness.{ApplyHarnessBinders, HarnessBinders, HarnessClockInstantiatorKey}
import chipyard.clocking.{SimplePllConfiguration, ClockDividerN}


// HarnessClockInstantiators are classes which generate clocks that drive
// TestHarness simulation models and any Clock inputs to the ChipTop
trait HarnessClockInstantiator {
  val clockMap: LinkedHashMap[String, (Double, Clock)] = LinkedHashMap.empty

  // request a clock at a particular frequency
  def requestClockHz(name: String, freqHzRequested: Double): Clock = {
    if (clockMap.contains(name)) {
      require(freqHzRequested == clockMap(name)._1,
        s"Request clock freq = $freqHzRequested != previously requested ${clockMap(name)._2} for requested clock $name")
      clockMap(name)._2
    } else {
      val clock = Wire(Clock())
      clockMap(name) = (freqHzRequested, clock)
      clock
    }
  }
  def requestClockMHz(name: String, freqMHzRequested: Double): Clock = {
    requestClockHz(name, freqMHzRequested * (1000 * 1000))
  }
  // refClock is the clock generated by TestDriver that is
  // passed to the TestHarness as its implicit clock
  def instantiateHarnessClocks(refClock: Clock, refClockFreqMHz: Double): Unit
}

class ClockSourceAtFreqMHz(val freqMHz: Double) extends BlackBox(Map(
  "PERIOD" -> DoubleParam(1000/freqMHz)
)) with HasBlackBoxInline {
  val io = IO(new ClockSourceIO)
  val moduleName = this.getClass.getSimpleName

  setInline(s"$moduleName.v",
    s"""
      |module $moduleName #(parameter PERIOD="") (
      |    input power,
      |    input gate,
      |    output clk);
      |  timeunit 1ns/1ps;
      |  reg clk_i = 1'b0;
      |  always #(PERIOD/2.0) clk_i = ~clk_i & (power & ~gate);
      |  assign clk = clk_i;
      |endmodule
      |""".stripMargin)
}


// The AbsoluteFreqHarnessClockInstantiator uses a Verilog blackbox to
// provide the precise requested frequency.
// This ClockInstantiator cannot be synthesized or run in FireSim
// It is useful for RTL simulations
class AbsoluteFreqHarnessClockInstantiator extends HarnessClockInstantiator {
  def instantiateHarnessClocks(refClock: Clock, refClockFreqMHz: Double): Unit = {
    // connect wires to clock source
    for ((name, (freqHz, clock)) <- clockMap) {
      val source = Module(new ClockSourceAtFreqMHz(freqHz / (1000 * 1000)))
      source.io.power := true.B
      source.io.gate := false.B

      clock := source.io.clk
    }
  }
}

class WithAbsoluteFreqHarnessClockInstantiator extends Config((site, here, up) => {
  case HarnessClockInstantiatorKey => () => new AbsoluteFreqHarnessClockInstantiator
})

class AllClocksFromHarnessClockInstantiator extends HarnessClockInstantiator {
  def instantiateHarnessClocks(refClock: Clock, refClockFreqMHz: Double): Unit = {
    val freqs = clockMap.map(_._2._1)
    freqs.tail.foreach(t => require(t == freqs.head, s"Mismatching clocks $t != ${freqs.head}"))
    for ((name, (freq, clock)) <- clockMap) {
      val freqMHz = freq / (1000 * 1000)
      require(freqMHz == refClockFreqMHz,
        s"AllClocksFromHarnessClockInstantiator has reference ${refClockFreqMHz.toInt} MHz attempting to drive clock $name which requires $freqMHz MHz")

      clock := refClock
    }
  }
}

class WithAllClocksFromHarnessClockInstantiator extends Config((site, here, up) => {
  case HarnessClockInstantiatorKey => () => new AllClocksFromHarnessClockInstantiator
})
