package wacc

import scala.collection.mutable
import cats.data.Chain
import wacc.ast.Position

private class LabelGenerator {
  import assemblyIR._
  import microWacc.{CallTarget, Ident, Builtin}
  import asmGenerator.escaped

  private val strings = mutable.HashMap[String, String]()
  private val files = mutable.HashMap[String, Int]()
  private var labelVal = -1
  private var permittedFuncFile: Option[String] = None

  /** Get an arbitrary label. */
  def getLabel(): String = {
    labelVal += 1
    s".L$labelVal"
  }

  private def getLabel(target: CallTarget | RuntimeError): String = target match {
    case Ident(v, guid)    => s"wacc_${v}_$guid"
    case Builtin(name)     => s"_$name"
    case err: RuntimeError => s".L.${err.name}"
  }

  /** Get a named label def for a function or error. */
  def getLabelDef(target: CallTarget | RuntimeError): LabelDef =
    LabelDef(getLabel(target))

  /** Get a named label for a function or error. */
  def getLabelArg(target: CallTarget | RuntimeError): LabelArg =
    LabelArg(getLabel(target))

  /** Get an arbitrary label for a string. */
  def getLabelArg(str: String): LabelArg =
    LabelArg(strings.getOrElseUpdate(str, s".L.str${strings.size}"))

  /** Get a named label for a string. */
  def getLabelArg(src: String, name: String): LabelArg =
    LabelArg(strings.getOrElseUpdate(src, s".L.$name.str${strings.size}"))

  /** Get a debug directive for a file. */
  def getDebugFile(file: java.io.File): Int =
    files.getOrElseUpdate(file.getCanonicalPath, files.size)

  /** Get a debug directive for a function. */
  def getDebugFunc(pos: Position, name: String, label: LabelDef): Chain[AsmLine] = {
    permittedFuncFile match {
      case Some(f) if f != pos.file.getCanonicalPath => Chain.empty
      case _ =>
        val customLabel = if name == "main" then Chain.empty else Chain(LabelDef(name))
        permittedFuncFile = Some(pos.file.getCanonicalPath)
        customLabel ++ Chain(
          Directive.Location(getDebugFile(pos.file), pos.line, None),
          Directive.Type(label, SymbolType.Function),
          Directive.Func(name, label)
        )
    }
  }

  /** Generate the assembly labels for constants that were labelled using the LabelGenerator. */
  def generateConstants: Chain[AsmLine] =
    strings.foldLeft(Chain.empty) { case (acc, (str, label)) =>
      acc ++ Chain(
        LabelDef(label),
        Directive.Asciz(str.escaped)
      )
    }

  /** Generates debug directives that were created using the LabelGenerator. */
  def generateDebug: Chain[AsmLine] =
    files.foldLeft(Chain.empty) { case (acc, (file, no)) =>
      acc :+ Directive.File(no, file)
    }
}