// SPDX-License-Identifier: Apache-2.0

package circt.stage.phases

import circt.Implicits.BooleanImplicits
import circt.stage.{CIRCTOptions, CIRCTTarget, EmittedMLIR, PreserveAggregate}

import firrtl.{
  AnnotationSeq,
  EmittedVerilogCircuit,
  EmittedVerilogCircuitAnnotation,
  EmittedVerilogModule,
  EmittedVerilogModuleAnnotation
}
import firrtl.annotations.{JsonProtocol, ReferenceTarget, Target}
import firrtl.options.{
  CustomFileEmission,
  Dependency,
  OptionsException,
  OutputAnnotationFileAnnotation,
  Phase,
  StageError,
  StageOptions,
  StageUtils
}
import firrtl.options.phases.WriteOutputAnnotations
import firrtl.options.Viewer.view
import firrtl.stage.{FirrtlOptions, RunFirrtlTransformAnnotation}
import _root_.logger.LogLevel

import java.io.File
import firrtl.annotations.LoadMemoryAnnotation

private object Helpers {
  implicit class LogLevelHelpers(logLevel: LogLevel.Value) {
    def toCIRCTOptions: Seq[String] = logLevel match {
      case LogLevel.Error => Seq.empty
      case LogLevel.Warn  => Seq.empty
      case LogLevel.Info  => Seq("-verbose-pass-executions")
      case LogLevel.Debug => Seq("-verbose-pass-executions")
      case LogLevel.Trace => Seq("-verbose-pass-executions", "-print-ir-after-all")
      case LogLevel.None  => Seq.empty
    }
  }

  /** Extract the JSON-encoded Annotation region from a single file output.
    *
    * @todo This is very janky and should be changed to something more stable.
    */
  def extractAnnotationFile(string: String, filename: String): (String, AnnotationSeq) = {
    var inAnno = false
    val (irLines, annoLines) = string.lines.foldLeft((Seq.empty[String], Seq.empty[String])) {
      case ((ir, anno), line) if line.startsWith("// ----- 8< ----- FILE") && line.contains(filename) =>
        inAnno = true
        (ir, anno)
      case ((ir, anno), line) if line.startsWith("// ----- 8< ----- FILE") =>
        inAnno = false
        (ir, anno)
      case ((ir, anno), line) if inAnno  => (ir, line +: anno)
      case ((ir, anno), line) if !inAnno => (line +: ir, anno)
    }
    val filtered = annoLines.reverse.toArray
      .mkString("\n")
    (
      irLines.reverse.mkString("\n"),
      filtered.forall(_.isWhitespace) match {
        case false => JsonProtocol.deserialize(filtered, false)
        case true  => Seq.empty
      }
    )
  }
}

private[this] object Exceptions {

  /** Wrap a message in colorful error text in the sytle of StageUtils.dramaticError.  That method prints to stdout and by
    * extracting this to just a method that does the string wrapping, this enables the message to be sent to another
    * file, e.g., stderr.
    * @todo Remove/unify this with StageUtils.dramaticError once SFC code is migrated into Chisel3.
    */
  def dramaticError(header: String, body: String): String = {
    s"""|$header
        |${"-" * 78}
        |$body
        |${"-" * 78}""".stripMargin
  }

  /** Indicates that the firtool binary failed with a non-zero exit code.  This generally indicates a compiler error
    * either originating from a user error or from a crash.
    *
    * @param binary the path to the firtool binary
    * @param exitCode the numeric exit status code returned by the firtool binary
    * @param stdout the contents returned to standard out by the firtool binary
    * @param stderr the contents returned to standard error by the firtool binary
    */
  class FirtoolNonZeroExitCode(binary: String, exitCode: Int, stdout: String, stderr: String)
      extends RuntimeException(
        dramaticError(
          header = s"${binary} returned a non-zero exit code",
          body = s"ExitCode:\n${exitCode}\nSTDOUT:\n${stdout}\nSTDERR:\n${stderr}"
        )
      )

  /** Indicates that the firtool binary was not found.  This likely indicates that the user didn't install
    * CIRCT/firtool.
    *
    * @param binary the path to the firtool binary
    */
  class FirtoolNotFound(binary: String)
      extends RuntimeException(
        dramaticError(
          header = s"$binary not found",
          body = """|Chisel requires that firtool, the MLIR-based FIRRTL Compiler (MFC), is installed
                    |and available on your $PATH.  (Did you forget to install it?)  You can download
                    |a binary release of firtool from the CIRCT releases webpage:
                    |  https://github.com/llvm/circt/releases""".stripMargin
        )
      )

}

/** A phase that calls and runs CIRCT, specifically `firtool`, while preserving an [[AnnotationSeq]] API.
  *
  * This is analogous to [[firrtl.stage.phases.Compiler]].
  */
class CIRCT extends Phase {

  import Helpers._
  import scala.sys.process._

  override def prerequisites = Seq(
    Dependency[firrtl.stage.phases.AddDefaults],
    Dependency[firrtl.stage.phases.AddImplicitEmitter],
    Dependency[firrtl.stage.phases.AddImplicitOutputFile]
  )
  override def optionalPrerequisites = Seq.empty
  override def optionalPrerequisiteOf = Seq.empty
  override def invalidates(a: Phase) = false

  override def transform(annotations: AnnotationSeq): AnnotationSeq = {
    val circtOptions = view[CIRCTOptions](annotations)
    val firrtlOptions = view[FirrtlOptions](annotations)
    val stageOptions = view[StageOptions](annotations)

    var blackbox, inferReadWrite = false
    var imcp = true
    var logLevel = _root_.logger.LogLevel.None
    var split = false
    val loadMems = collection.mutable.ArrayBuffer[LoadMemoryAnnotation]()

    val annotationsx: AnnotationSeq = annotations.flatMap {
      case a: CustomFileEmission => {
        val filename = a.filename(annotations)
        a.replacements(filename)
      }
      case _: firrtl.EmitCircuitAnnotation => Nil
      case _: firrtl.EmitAllModulesAnnotation => {
        split = true
        Nil
      }
      case a @ RunFirrtlTransformAnnotation(transform) =>
        transform match {
          /* Inlining/Flattening happen by default, so these can be dropped. */
          case _: firrtl.passes.InlineInstances | _: firrtl.transforms.Flatten => Nil
          /* ReplSeqMem is converted to a firtool option */
          case _: firrtl.passes.memlib.ReplSeqMem =>
            blackbox = true
            Nil
          /* Any emitters should not be passed to firtool. */
          case _: firrtl.Emitter => Nil
          /* Default case: leave the annotation around and let firtool warn about it. */
          case _ => Seq(a)
        }
      case firrtl.passes.memlib.InferReadWriteAnnotation =>
        inferReadWrite = true
        Nil
      case firrtl.transforms.NoConstantPropagationAnnotation =>
        imcp = false
        Nil
      case anno: _root_.logger.LogLevelAnnotation =>
        logLevel = anno.globalLogLevel
        Nil
      /* The following can be dropped. */
      case _: firrtl.transforms.CombinationalPath   => Nil
      case _: _root_.logger.ClassLogLevelAnnotation => Nil
      /* Specially handled */
      case l: firrtl.annotations.LoadMemoryAnnotation =>
        loadMems += l
        Seq(chisel3.experimental.Trace.traceName(l.getTargets.head))
      /* Default case: leave the annotation around and let firtool warn about it. */
      case a => Seq(a)
    }

    /* Filter the annotations to only those things which CIRCT should see. */
    (new WriteOutputAnnotations).transform(annotationsx)

    val input: String = firrtlOptions.firrtlCircuit match {
      case None          => throw new OptionsException("No input file specified!")
      case Some(circuit) => circuit.serialize
    }

    val chiselAnnotationFileName: Option[String] =
      stageOptions.annotationFileOut.map(stageOptions.getBuildFileName(_, Some(".anno.json")))

    val circtAnnotationFileName = "circt.anno.json"

    val binary = "firtool"

    val cmd =
      Seq(binary, "-format=fir", "-warn-on-unprocessed-annotations", "-dedup") ++
        Seq("-output-annotation-file", circtAnnotationFileName) ++
        circtOptions.firtoolOptions ++
        logLevel.toCIRCTOptions ++
        /* The following options are on by default, so we disable them if they are false. */
        (circtOptions.preserveAggregate match {
          case Some(PreserveAggregate.OneDimVec) => Seq("-preserve-aggregate=1d-vec")
          case Some(PreserveAggregate.Vec)       => Seq("-preserve-aggregate=vec")
          case Some(PreserveAggregate.All)       => Seq("-preserve-aggregate=all")
          case None                              => None
        }) ++
        circtOptions.preserveAggregate.map(_ => "-preserve-public-types=0") ++
        (!inferReadWrite).option("-disable-infer-rw") ++
        (!imcp).option("-disable-imcp") ++
        /* The following options are off by default, so we enable them if they are true. */
        (blackbox).option("-blackbox-memory") ++
        /* Communicate the annotation file through a file. */
        (chiselAnnotationFileName.map(a => Seq("-annotation-file", a))).getOrElse(Seq.empty) ++
        /* Convert the target to a firtool-compatible option. */
        ((circtOptions.target, split) match {
          case (Some(CIRCTTarget.FIRRTL), false)        => Seq("-ir-fir")
          case (Some(CIRCTTarget.HW), false)            => Seq("-ir-hw")
          case (Some(CIRCTTarget.Verilog), true)        => Seq("--split-verilog", s"-o=${stageOptions.targetDir}")
          case (Some(CIRCTTarget.Verilog), false)       => None
          case (Some(CIRCTTarget.SystemVerilog), true)  => Seq("--split-verilog", s"-o=${stageOptions.targetDir}")
          case (Some(CIRCTTarget.SystemVerilog), false) => None
          case (None, _) =>
            throw new Exception(
              "No 'circtOptions.target' specified. This should be impossible if dependencies are satisfied!"
            )
          case (_, true) =>
            throw new Exception(
              s"The circtOptions.target specified (${circtOptions.target}) does not support running with an EmitAllModulesAnnotation as CIRCT only supports one-file-per-module for Verilog or SystemVerilog targets."
            )
          case _ =>
            throw new Exception(
              s"Invalid combination of circtOptions.target ${circtOptions.target} and split ${split}"
            )
        })

    logger.info(s"""Running CIRCT: '${cmd.mkString(" ")} < $$input'""")
    println(s"""Running CIRCT: '${cmd.mkString(" ")} < $$input'""")
    val stdoutStream, stderrStream = new java.io.ByteArrayOutputStream
    val stdoutWriter = new java.io.PrintWriter(stdoutStream)
    val stderrWriter = new java.io.PrintWriter(stderrStream)
    val exitValue =
      try {
        (cmd #< new java.io.ByteArrayInputStream(input.getBytes))
          .!(ProcessLogger(stdoutWriter.println, stderrWriter.println))
      } catch {
        case a: java.lang.RuntimeException if a.getMessage().startsWith("No exit code") =>
          throw new Exceptions.FirtoolNotFound(binary)
      }
    stdoutWriter.close()
    stderrWriter.close()
    val result = stdoutStream.toString
    logger.info(result)
    val errors = stderrStream.toString
    if (exitValue != 0)
      throw new Exceptions.FirtoolNonZeroExitCode(binary, exitValue, result, errors)
    val circtOutputAnnotations = if (split) {
      logger.info(result)
      val file = new File(stageOptions.getBuildFileName(circtAnnotationFileName, Some(".anno.json")))
      file match {
        case file if !file.canRead() => Seq.empty
        case file => {
          val foo = JsonProtocol.deserialize(file, false)
          foo
        }
      }
    } else { // if split it has already been written out to the file system stdout not necessary.
      val outputFileName: String = stageOptions.getBuildFileName(firrtlOptions.outputFileName.get)
      val (output, outputAnnotations) = extractAnnotationFile(result, circtAnnotationFileName)
      outputAnnotations ++ (circtOptions.target match {
        case Some(CIRCTTarget.FIRRTL) =>
          Seq(EmittedMLIR(outputFileName, output, Some(".fir.mlir")))
        case Some(CIRCTTarget.HW) =>
          Seq(EmittedMLIR(outputFileName, output, Some(".hw.mlir")))
        case Some(CIRCTTarget.Verilog) =>
          Seq(EmittedVerilogCircuitAnnotation(EmittedVerilogCircuit(outputFileName, output, ".v")))
        case Some(CIRCTTarget.SystemVerilog) =>
          Seq(EmittedVerilogCircuitAnnotation(EmittedVerilogCircuit(outputFileName, output, ".sv")))
        case None =>
          throw new Exception(
            "No 'circtOptions.target' specified. This should be impossible if dependencies are satisfied!"
          )

      })
    }

    // -------- Handle loadMemoryFromFile annotation ----------
    // TODO: develop a better first-class feature in MFC
    val mapper = chisel3.experimental.Trace.finalTargetMap(circtOutputAnnotations)
    var memoryCounter = -1
    val allMemoryReferences = loadMems.flatMap { lm: LoadMemoryAnnotation =>
      mapper(lm.getTargets.head).map {
        case x: ReferenceTarget => lm.copy(target = x.encapsulatingModuleTarget.ref(x.ref).toNamed)
      }
    }.distinct // updated for deduplication
    val boundMemoryFileAnnotationAndFileName = allMemoryReferences
      .groupBy(_.target.module.name)
      .toList
      .sortBy(_._1) // Giving a deterministic iteration order for memoryCounter
      .flatMap {
        case (m: String, lms: Seq[LoadMemoryAnnotation]) =>
          lms.sortBy(_.target.serialize) // Giving a deterministic iteration order for memoryCounter
             .zipWithIndex.map {
            case (lm: LoadMemoryAnnotation, i: Int) =>
              memoryCounter += 1
              val contents = chisel3.util.experimental.loadMemoryFromFile.buildVerilog(
                lm.fileName,
                lm.target,
                memoryCounter,
                lm.hexOrBinary.serialize
              )

              val name = s"${lm.target.serialize}"
              val outputFileNameWithoutSuffix = stageOptions.getBuildFileName(name)
              (
                EmittedVerilogModuleAnnotation(new EmittedVerilogModule(outputFileNameWithoutSuffix, contents, ".v")),
                outputFileNameWithoutSuffix
              )
          }
      }
      .toList

    // Re-using the EmittedVerilogModuleAnnotation to write a ".f" file at the end of Stage/Phase serialization
    val dotFFileName = (new File(stageOptions.getBuildFileName(chisel3.util.experimental.loadMemoryFromFile.defaultFileListName))).getAbsolutePath()
    val dotFFileAnnotation = EmittedVerilogModuleAnnotation(
      new EmittedVerilogModule(
        dotFFileName,
        boundMemoryFileAnnotationAndFileName.map { case (_, o) => new File(o + ".v").getAbsoluteFile() }.mkString("\n"),
        ".f"
      )
    )
    val boundMemoryFileAnnotations = boundMemoryFileAnnotationAndFileName.map(_._1)
    // -------- Handled loadMemoryFromFile annotation ----------

    dotFFileAnnotation +: (boundMemoryFileAnnotations ++ circtOutputAnnotations)
  }

}
