package uvm.refimpl.bootimg

import java.io._
import java.io.OutputStream
import java.nio.charset.StandardCharsets
import java.nio.file._
import java.util.zip.ZipEntry
import java.util.zip.ZipOutputStream

import scala.collection.mutable.ArrayBuffer
import scala.collection.mutable.HashMap

import org.slf4j.LoggerFactory

import com.typesafe.scalalogging.Logger

import uvm._
import uvm.ir.irbuilder.IRBuilder
import uvm.ir.textoutput.BundleSerializer
import uvm.ir.textoutput.EntityUtils
import uvm.refimpl._
import uvm.refimpl.itpr.FrameCursor
import uvm.refimpl.itpr.InterpreterStack
import uvm.refimpl.itpr.InterpreterThread
import uvm.refimpl.mem.HeaderUtils
import uvm.refimpl.mem.Space
import uvm.refimpl.mem.TypeSizes
import uvm.refimpl.mem.scanning.MemoryDataScanner
import uvm.refimpl.mem.scanning.MemoryFieldHandler
import uvm.refimpl.nat.NativeSupport
import uvm.types._
import uvm.utils.IOHelpers
import uvm.utils.WithUtils.tryWithResource

class BootImageBuilder(implicit microVM: MicroVM) {
  def makeBootImage(whiteList: Seq[TopLevel], primordial: PrimordialInfo,
    syms: Seq[FieldAndSymbol], relocs: Seq[FieldAndSymbol], outputFile: String): Unit = {
    implicit val tcb = new TransitiveClosureBuilder(whiteList, primordial)

    tcb.doTransitiveClosure()

    val biw = new BootImageWriter(tcb, syms, relocs, outputFile)
    biw.writeFile()
  }
}

case class PrimordialInfo(maybeFunc: Option[Function], maybeStack: Option[InterpreterStack], maybeThreadLocal: Option[Word])
object PrimordialInfo {
  val NO_PRIMORDIAL = PrimordialInfo(None, None, None)
}

case class FieldAndSymbol(objRef: Word, offset: Word, symbol: String) {
  def toAddrSymPair: (Word, String) = ((objRef + offset), symbol)
}

object BootImageWriter {
  import BootImageFile._
  val logger = Logger(LoggerFactory.getLogger(getClass.getName))

  class FileGroup(baseName: String, dir: Path) {
    val dataFile = dir.resolve(baseName + FILEEXT_DATA)
    val allocFile = dir.resolve(baseName + FILEEXT_ALLOC)
    val relocFile = dir.resolve(baseName + FILEEXT_RELOC)

    val dataOS = new AlignedOutputStream(new BufferedOutputStream(Files.newOutputStream(dataFile)))

    val allocRecs = new ArrayBuffer[UnitAllocRecord]
    val relocRecs = new ArrayBuffer[FieldRelocRecord]
  }
}

class BootImageWriter(tcb: TransitiveClosureBuilder, syms: Seq[FieldAndSymbol],
    relocs: Seq[FieldAndSymbol], outputFile: String)(implicit microVM: MicroVM) {
  import BootImageWriter._
  import BootImageFile._

  val bundleSerializer = new BundleSerializer(microVM.globalBundle, tcb.tls.set)

  val tempDir = Files.createTempDirectory("mu-bootimg")

  logger.info("Writing bootimg files to temp dir: %s".format(tempDir.toAbsolutePath().toString()))

  val globalGroup = new FileGroup("global", tempDir)
  val heapGroup = new FileGroup("heap", tempDir)

  val manualSymRecs = new ArrayBuffer[ManualSymRecord]

  val relocsMap = relocs.map(_.toAddrSymPair).toMap

  {
    logger.debug("Relocation map: " + relocsMap.toString)
  }

  val idNameMapPath = tempDir.resolve(IDNAMEMAP_FILE)
  val uirBundlePath = tempDir.resolve(UIRBUNDLE_FILE)
  val manualSymsPath = tempDir.resolve(MANUALSYMS_FILE)
  val metaInfoPath = tempDir.resolve(METAINFO_FILE)
  val zipPath = Paths.get(outputFile)

  val allZipContents = Seq(globalGroup.dataFile, heapGroup.dataFile, globalGroup.allocFile, heapGroup.allocFile,
    globalGroup.relocFile, heapGroup.relocFile, idNameMapPath, uirBundlePath, manualSymsPath, metaInfoPath)

  private val globalMemory: Space = microVM.memoryManager.globalMemory
  private val smallObjectSpace: Space = microVM.memoryManager.heap.space
  private val largeObjectSpace: Space = microVM.memoryManager.heap.los
  private implicit val memorySupport = microVM.memoryManager.memorySupport

  val addrToGlobalCell = tcb.globalCellMap
  val addrToHeapObjNum = new HashMap[Word, Long]()

  def writeFile(): Unit = {
    for (addr <- tcb.allocs.set) {
      if (globalMemory.isInSpace(addr)) {
        writeAllocUnit(addr, globalGroup, isGlobal = true)
      } else {
        writeAllocUnit(addr, heapGroup, isGlobal = false)
      }
    }
    globalGroup.dataOS.close()
    heapGroup.dataOS.close()

    writeAllocRecs(globalGroup)
    writeAllocRecs(heapGroup)

    locateManualSyms()
    writeManualSyms()

    for (ar <- heapGroup.allocRecs) {
      addrToHeapObjNum(ar.addr) = ar.num
    }

    scanForRelocs(globalGroup)
    scanForRelocs(heapGroup)

    addPointerRelocs()

    writeRelocRecs(globalGroup)
    writeRelocRecs(heapGroup)

    writeIDNameMap()
    writeUIRBundle()

    writeMetaInfo()

    makeZipPackage()
  }

  private def writeAllocUnit(addr: Word, group: FileGroup, isGlobal: Boolean): Unit = {
    val tag = HeaderUtils.getTag(addr)
    val ty = HeaderUtils.getType(microVM, tag)
    val (size, align, varLen) = ty match {
      case t: TypeHybrid => {
        val len = HeaderUtils.getVarLength(addr)
        val sz = TypeSizes.hybridSizeOf(t, len)
        val al = TypeSizes.hybridAlignOf(t, len)
        (sz, al, len)
      }
      case t => {
        val sz = TypeSizes.sizeOf(t)
        val al = TypeSizes.alignOf(t)
        (sz, al, 0L)
      }
    }

    val num = if (isGlobal) addrToGlobalCell(addr).g.id else group.allocRecs.length

    group.dataOS.alignUp(align)
    val pos = group.dataOS.position
    group.dataOS.writeFromMemory(addr, size)

    val rec = UnitAllocRecord(num, pos, ty, varLen, addr)
    group.allocRecs += rec
  }

  private def writeAllocRecs(group: FileGroup): Unit = {
    tryWithResource(Files.newBufferedWriter(group.allocFile, StandardCharsets.UTF_8)) { writer =>
      for (ar <- group.allocRecs) {
        writer.write("%d,%d,%d,%d,%s\n".format(ar.num, ar.fileOffset, ar.ty.id, ar.varLen, EntityUtils.getName(ar.ty)))
      }
    }
  }

  private def toManualSymRec(fns: FieldAndSymbol): ManualSymRecord = {
    val iref = fns.objRef + fns.offset
    if (fns.objRef != 0L || !globalMemory.isInSpace(iref)) {
      throw new BootImageBuilderException("Symbol can only be in global cells. Symbol: '%s', objRef: 0x%x, offset: 0x%x, addr: 0x%x".format(
        fns.symbol, fns.objRef, fns.offset, iref))
    }

    val gcr = tcb.getGlobalCellRec(iref)
    val gid = gcr.g.id
    val offset = fns.offset - gcr.begin

    val msr = ManualSymRecord(fns.symbol, gid, offset)

    msr
  }

  private def locateManualSyms(): Unit = {
    for (fns <- syms) {
      manualSymRecs += toManualSymRec(fns)
    }
  }

  private def writeManualSyms(): Unit = {
    tryWithResource(Files.newBufferedWriter(manualSymsPath, StandardCharsets.UTF_8)) { writer =>
      for (msr <- manualSymRecs) {
        writer.write("%s,%d,%d\n".format(
          msr.sym, msr.targetNum, msr.targetOffset))
      }
    }
  }

  private def scanForRelocs(group: FileGroup): Unit = {
    for (alloc <- group.allocRecs) {
      val num = alloc.num
      val begin = alloc.addr

      MemoryDataScanner.scanAllocUnit(begin, new MemoryFieldHandler {
        def visitUPtrField(objRef: Word, iRef: Word, toAddr: Word): Unit = {
          relocsMap.get(iRef) match {
            case Some(symbol) => {
              // by symbol
              val fieldOffset = iRef - objRef
              val reloc = FieldRelocRecord(num, fieldOffset, AS_UPTR, TO_SYM, 0, 0, symbol)
              group.relocRecs += reloc
              logger.debug("External uptr relocation: %d 0x%x -> %s".format(iRef, iRef, symbol))
            }
            case None => {
              // maybe a "reasonable" pointer to a mu global cell. Relocate it too.
              val fieldOffset = iRef - objRef
              val maybeGlobalCell = tcb.maybeGetGlobalCellRec(toAddr)
              maybeGlobalCell match {
                case Some(gcr) =>
                  val targetNum = gcr.g.id.toLong
                  val targetAddr = gcr.begin
                  val targetOffset = toAddr - targetAddr
                  val reloc = FieldRelocRecord(num, fieldOffset, AS_UPTR, TO_GLOBAL, targetNum, targetOffset, NO_SYM)
                  group.relocRecs += reloc
                  logger.info("Relocation entry for uptr field automatically generated: %d 0x%x -> %d 0x%x".format(
                    iRef, iRef, toAddr, toAddr))
                case None =>
              }
            }
          }
        }
        def visitUFPField(objRef: Word, iRef: Word, toAddr: Word): Unit = {
          relocsMap.get(iRef).foreach { symbol =>
            // by symbol
            val fieldOffset = iRef - objRef
            val reloc = FieldRelocRecord(num, fieldOffset, AS_UPTR, TO_SYM, 0, 0, symbol)
            group.relocRecs += reloc
            logger.debug("External ufuncptr relocation: %d 0x%x -> %s".format(iRef, iRef, symbol))
          }
        }

        def visitRefField(objRef: Word, iRef: Word, toObj: Word, isWeak: Boolean): Unit = if (toObj != 0L) {
          val fieldOffset = iRef - objRef
          val targetNum = addrToHeapObjNum(toObj)
          val reloc = FieldRelocRecord(num, fieldOffset, AS_REF, TO_HEAP, targetNum, NO_OFFSET, NO_SYM)
          group.relocRecs += reloc
        }

        def visitIRefField(objRef: Word, iRef: Word, toObj: Word, toOffset: Word): Unit = if (toObj != 0L) {
          val fieldOffset = iRef - objRef
          val targetNum = addrToHeapObjNum(toObj)
          val reloc = FieldRelocRecord(num, fieldOffset, AS_IREF, TO_HEAP, targetNum, toOffset, NO_SYM)
          group.relocRecs += reloc
        } else {
          val fieldOffset = iRef - objRef
          val gcr = tcb.getGlobalCellRec(toOffset)
          val targetNum = gcr.g.id.toLong
          val targetAddr = gcr.begin
          val targetOffset = toOffset - targetAddr
          val reloc = FieldRelocRecord(num, fieldOffset, AS_IREF, TO_GLOBAL, targetNum, targetOffset, NO_SYM)
          group.relocRecs += reloc
        }

        def visitTagRefField(objRef: Word, iRef: Word, toObj: Word): Unit = if (toObj != 0L) {
          val fieldOffset = iRef - objRef
          val targetNum = addrToHeapObjNum(toObj)
          val reloc = FieldRelocRecord(num, fieldOffset, AS_TAGREF, TO_HEAP, targetNum, NO_OFFSET, NO_SYM)
          group.relocRecs += reloc
        }

        def visitFuncRefField(objRef: Word, iRef: Word, toFunc: Option[Function]): Unit = toFunc.foreach { func =>
          val fieldOffset = iRef - objRef
          val targetNum = func.id.toLong
          val reloc = FieldRelocRecord(num, fieldOffset, AS_FUNCREF, TO_FUNC, targetNum, NO_OFFSET, NO_SYM)
          group.relocRecs += reloc
        }

        // We already made sure in the transitive closure step that the following are all NULL.
        def visitStackRefField(objRef: Word, iRef: Word, toStack: Option[InterpreterStack]): Unit = {}
        def visitThreadRefField(objRef: Word, iRef: Word, toThread: Option[InterpreterThread]): Unit = {}
        def visitFCRefField(objRef: Word, iRef: Word, toFCRef: Option[FrameCursor]): Unit = {}
        def visitIRBuilderRefField(objRef: Word, iRef: Word, toIRNode: Option[IRBuilder]): Unit = {}
      })
    }
  }

  private def addPointerRelocs(): Unit = {
    // Seems uptrs are already handled when scanning fields. No further processing is needed at this moment.
  }

  private def writeRelocRecs(group: FileGroup): Unit = {
    tryWithResource(Files.newBufferedWriter(group.relocFile, StandardCharsets.UTF_8)) { writer =>
      for (rr <- group.relocRecs) {
        writer.write("%d,%d,%s,%s,%d,%d,%s\n".format(
          rr.num, rr.fieldOffset, rr.fieldKind, rr.targetKind, rr.targetNum, rr.targetOffset, rr.targetString))
      }
    }
  }

  private def writeIDNameMap(): Unit = {
    tryWithResource(Files.newBufferedWriter(idNameMapPath, StandardCharsets.UTF_8)) { writer =>
      bundleSerializer.writeIDNameMap(writer)
    }
  }

  private def writeUIRBundle(): Unit = {
    tryWithResource(Files.newBufferedWriter(uirBundlePath, StandardCharsets.UTF_8)) { writer =>
      bundleSerializer.writeUIR(writer)
    }
  }

  private def writeMetaInfo(): Unit = {
    tryWithResource(Files.newBufferedWriter(metaInfoPath, StandardCharsets.UTF_8)) { writer =>
      val sb = new StringBuilder()
      tcb.primordial.maybeFunc.foreach { func =>
        sb.append("%s=%d\n".format(KEY_INITFUNC, func.id))
        tcb.primordial.maybeThreadLocal.foreach { addr =>
          val heapObjNum = heapGroup.allocRecs.find(_.addr == addr).getOrElse {
            throw new BootImageBuilderException("BUG: Thread local object is not in the boot image")
          }.num
          sb.append("%s=%d\n".format(KEY_THREADLOCAL, heapObjNum))
        }
      }
      writer.write(sb.toString)
    }
  }

  private def makeZipPackage(): Unit = {
    tryWithResource(new ZipOutputStream(Files.newOutputStream(zipPath))) { zos =>
      for (path <- allZipContents) {
        val entry = new ZipEntry(path.getFileName.toString)
        zos.putNextEntry(entry)
        tryWithResource(Files.newInputStream(path)) { fis =>
          IOHelpers.copy(fis, zos)
        }
      }
    }
  }
}

class AlignedOutputStream(os: OutputStream) extends Closeable {
  var position: Long = 0L

  private val BUFFER_SZ = 16384
  private val buffer = new Array[Byte](BUFFER_SZ)

  def alignUp(alignment: Long): Unit = {
    val newPos = TypeSizes.alignUp(position, alignment)

    position until newPos foreach { _ => os.write(0) }

    position = newPos
  }

  def writeFromMemory(addr: Long, size: Long): Unit = {
    var nextAddr = addr
    var remain = size

    while (remain > BUFFER_SZ) {
      transfer(nextAddr, BUFFER_SZ)
      nextAddr += BUFFER_SZ
      remain -= BUFFER_SZ
    }

    if (remain > 0) {
      transfer(nextAddr, remain.toInt)
    }

    position += size
  }

  private def transfer(addr: Long, size: Int): Unit = {
    NativeSupport.theMemory.get(addr, buffer, 0, size)
    os.write(buffer, 0, size)
  }

  def close(): Unit = {
    os.close()
  }
}