import firrtl._
import firrtl.annotations.{Annotation, SingleTargetAnnotation}
import firrtl.annotations.{CircuitTarget, ModuleTarget, InstanceTarget, ReferenceTarget, Target}

import chisel3._
import chisel3.experimental.{annotate, ChiselAnnotation, RunFirrtlTransform}
import chisel3.internal.InstanceId

case class InfoAnnotation(target: Target, info: String) extends SingleTargetAnnotation[Target] {
  def duplicate(newTarget: Target) = this.copy(target = newTarget)
}
object InfoAnnotator {
  def info(component: InstanceId, info: String): Unit = {
    annotate(new ChiselAnnotation with RunFirrtlTransform {
      def toFirrtl: Annotation = InfoAnnotation(component.toTarget, info)
      def transformClass = classOf[InfoTransform]
    })
  }
}

/** A transform that reads InfoAnnotations and prints information about them */
class InfoTransform() extends Transform with DependencyAPIMigration {

  override def prerequisites = firrtl.stage.Forms.HighForm

  override def execute(state: CircuitState): CircuitState = {
    println("Starting transform 'IdentityTransform'")

    // Add infoanno - a single format for all object types
    val annotationsx = state.annotations.flatMap{
      case InfoAnnotation(a: CircuitTarget, info) | InfoAnnotation(a: ModuleTarget, info) | InfoAnnotation(a: InstanceTarget, info) | InfoAnnotation(a: ReferenceTarget, info) =>
        println(s"PROOF_MARKER {${a.serialize}, '$info'}")
        None
      case a =>
        Some(a)
    }

    state.copy(annotations = annotationsx)
  }
}


class RegisterFile() extends Module {
  val io = IO(new Bundle {
    val readreg1  = Input(UInt(5.W))
    val readreg2  = Input(UInt(5.W))
    val writereg  = Input(UInt(5.W))
    val writedata = Input(UInt(32.W))
    val wen       = Input(Bool())

    val readdata1 = Output(UInt(32.W))
    val readdata2 = Output(UInt(32.W))
  })

  // Required so the compiler doesn't optimize things away when testing
  // incomplete designs.
  dontTouch(io)
  val regs = Reg(Vec(32, UInt(32.W)))
  for (i <- 0 to regs.length-1) {
    InfoAnnotator.info(regs(i), s"x$i")
  }
  // When the write enable is high, write the data
  when (io.wen) {
    regs(io.writereg) := io.writedata
  }

  // *Always* read the data. This is required for the single cycle CPU since in a single cycle it
  // might both read and write the registers (e.g., an add)
  io.readdata1 := regs(io.readreg1)
  io.readdata2 := regs(io.readreg2)

}

println(getVerilog(new ALU))

// ChiselStage.emitChirrtl(new RegisterFile())
import chisel3.stage.{ChiselStage, ChiselGeneratorAnnotation}
// TODO fix type error here
  (new ChiselStage).execute(Array.empty, Seq(ChiselGeneratorAnnotation(() => new RegisterFile())))