package uvm.refimpl.bootimg

import scala.collection.immutable.TreeMap
import scala.collection.mutable.{ Set, Queue }

import org.slf4j.LoggerFactory

import com.typesafe.scalalogging.Logger

import uvm._
import uvm.ir.irbuilder.IRBuilder
import uvm.refimpl._
import uvm.refimpl.itpr._
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.ssavariables._
import uvm.types._

object TransitiveClosure {
  def apply[T](initialElems: T*) = new TransitiveClosure(initialElems)
}

class TransitiveClosure[T](initialElems: Seq[T]) {
  val set = Set[T](initialElems: _*)
  val queue = Queue[T](initialElems: _*)

  def hasPending = !queue.isEmpty
  def dequeue() = queue.dequeue()
  def maybeEnqueue(elem: T): Boolean = {
    if (set.contains(elem)) {
      false
    } else {
      set += elem
      queue += elem
      true
    }
  }
  def +=(elem: T): TransitiveClosure[T] = { maybeEnqueue(elem); this }
  def ++=(elems: Seq[T]): TransitiveClosure[T] = { elems.foreach(this +=); this }
}

object TransitiveClosureBuilder {
  val logger = Logger(LoggerFactory.getLogger(getClass.getName))

  implicit class RichTransitiveClosure[T <: Identified](val tc: TransitiveClosure[T]) extends AnyVal {
    def ?=(elem: SSAVariable): TransitiveClosure[T] = {
      if (elem.isInstanceOf[GlobalVariable]) {
        tc += elem.asInstanceOf[T]
      }
      tc
    }
    def ??=(elems: Seq[SSAVariable]): TransitiveClosure[T] = {
      for (elem <- elems) {
        this ?= elem
      }
      tc
    }
    def ??=(dest: DestClause): TransitiveClosure[T] = {
      this ??= dest.args
      tc
    }
    def ??=(maybeExcClause: Option[ExcClause]): TransitiveClosure[T] = {
      maybeExcClause foreach { excClause =>
        this ??= excClause.nor.args ??= excClause.exc.args
      }
      tc
    }
    def ???=(maybeElem: Option[SSAVariable]): TransitiveClosure[T] = { // More ? to work around type erasure
      for (elem <- maybeElem) {
        this ?= elem
      }
      tc
    }
    def ????=(maybeDest: Option[DestClause]): TransitiveClosure[T] = {
      for (dest <- maybeDest) {
        this ??= dest.args
      }
      tc
    }
  }

  case class GlobalCellRec(begin: Word, end: Word, g: GlobalCell)
}

class TransitiveClosureBuilder(initialSet: Seq[TopLevel], val primordial: PrimordialInfo)(
    implicit microVM: MicroVM) extends MemoryFieldHandler {
  import TransitiveClosureBuilder._

  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

  /**
   * Top-level definitions in the closure. Function versions are not considered top-level here. A function may have at
   * most one function version. Only that version is serialised into the boot image.
   */
  val tls = TransitiveClosure[TopLevel](initialSet: _*)

  /**
   * Global cells + heap objects.
   */
  val allocs = TransitiveClosure[Word]()
  
  // Put the primordial function/stack and the thread-local into the set
  {
    val PrimordialInfo(pf, ps, ptl) = primordial
    
    if (ps.isDefined) {
      throw new BootImageBuilderException("Building boot image using a stack is not implemented. The current way is to specify the stack-bottom function.")
    }
    
    pf foreach { func =>
      // If the primordial function/stack is specified, add it.
      tls += func
      
      ptl foreach { addr =>
        // Also add the thread-local if specified.
        allocs += addr
      }
    }
  }

  private def getGlobalCellAddr(g: GlobalCell): Word = {
    microVM.constantPool.getGlobalVarBox(g).asInstanceOf[BoxIRef].addr
  }

  /**
   * Map the global cell's starting address to the global cell itself.
   */
  val globalCellMap = {
    val elems = microVM.globalBundle.globalCellNs.all.toSeq.map { g =>
      val begin = getGlobalCellAddr(g)
      val ty = g.cellTy
      assert(!ty.isInstanceOf[TypeHybrid])
      val size = TypeSizes.sizeOf(ty)
      val end = begin + size
      (begin, GlobalCellRec(begin, end, g))
    }
    TreeMap[Word, GlobalCellRec](elems: _*)
  }

  def getGlobalCellRecExact(begin: Word): GlobalCellRec = {
    require(microVM.memoryManager.globalMemory.isInSpace(begin),
      "Address %d 0x%x is not in the global space".format(begin, begin))

    val gcr = globalCellMap(begin)

    gcr
  }

  def getGlobalCellRec(iref: Word): GlobalCellRec = {
    require(microVM.memoryManager.globalMemory.isInSpace(iref),
      "Address %d 0x%x is not in the global space".format(iref, iref))

    val (begin, gcr) = globalCellMap.to(iref).last
    assert(begin <= iref && iref < gcr.end,
      "Address %d 0x%x is in the global space, but the previous global cell's address range is %d 0x%x - %d 0x%x".format(
        iref, iref, begin, begin, gcr.end, gcr.end))

    gcr
  }

  def maybeGetGlobalCellRec(iref: Word): Option[GlobalCellRec] = {
    if (microVM.memoryManager.globalMemory.isInSpace(iref)) {
      for ((begin, gcr) <- globalCellMap.to(iref).lastOption) yield {
        assert(begin <= iref && iref < gcr.end,
          "Address %d 0x%x is in the global space, but the previous global cell's address range is %d 0x%x - %d 0x%x".format(
            iref, iref, begin, begin, gcr.end, gcr.end))

        gcr
      }
    } else { 
      None
    }
  }

  def doTransitiveClosure(): Unit = {
    while (tls.hasPending || allocs.hasPending) {
      while (tls.hasPending) {
        visitTopLevel(tls.dequeue())
      }
      while (allocs.hasPending) {
        visitAllocUnit(allocs.dequeue())
      }
    }
  }

  private def visitTopLevel(tl: Identified): Unit = {
    tl match {
      case t: Type        => visitType(t)
      case s: FuncSig     => visitFuncSig(s)
      case c: Constant    => visitConstant(c)
      case g: GlobalCell  => visitGlobalCell(g)
      case f: Function    => visitFunction(f)
      case e: ExposedFunc => visitExpFunc(e)
    }
  }

  private def visitType(ty: Type): Unit = ty match {
    case TypeUPtr(t)       => tls += t
    case TypeUFuncPtr(s)   => tls += s
    case TypeStruct(ts)    => tls ++= ts
    case TypeHybrid(fs, v) => tls ++= fs += v
    case TypeArray(t, _)   => tls += t
    case TypeVector(t, _)  => tls += t
    case TypeRef(t)        => tls += t
    case TypeIRef(t)       => tls += t
    case TypeWeakRef(t)    => tls += t
    case TypeFuncRef(s)    => tls += s
    case _                 => // Not nested or recursive. Do nothing.
  }

  private def visitFuncSig(s: FuncSig): Unit = {
    tls ++= s.paramTys ++= s.retTys
  }

  private def visitConstant(c: Constant): Unit = {
    tls += c.constTy

    c match {
      case ConstSeq(_, elems) => tls ++= elems
      case _                  => // Not nested. Do nothing.
    }
  }

  private def visitGlobalCell(g: GlobalCell): Unit = {
    tls += g.cellTy

    allocs += getGlobalCellAddr(g)
  }

  private def visitFunction(f: Function): Unit = {
    tls += f.sig

    microVM.globalBundle.funcToVers(f).headOption match {
      case None => // Undefined. Do nothing
      case Some(mostRecent) => {
        // Visit all TOP-LEVELS (mostly types sigs and SSA variables).
        // This version, including all of its basic blocks and instructions, will be in the boot image. 
        for (bb <- mostRecent.bbs) {
          tls ++= bb.norParams.map(_.ty)

          for (inst <- bb.insts) {
            visitInstruction(inst)
          }
        }
      }
    }
  }

  private def visitExpFunc(e: ExposedFunc): Unit = {
    tls += e.func
    tls += e.cookie
  }

  private def visitInstruction(inst: Instruction): Unit = inst match {
    case InstBinOp(op, opndTy, op1, op2, excClause)              => tls += opndTy ?= op1 ?= op2 ??= excClause
    case InstCmp(op, opndTy, op1, op2)                           => tls += opndTy ?= op1 ?= op2
    case InstConv(op, fromTy, toTy, opnd)                        => tls += fromTy += toTy ?= opnd
    case InstSelect(condTy, opndTy, cond, ifTrue, ifFalse)       => tls += condTy += opndTy ?= cond ?= ifTrue ?= ifFalse
    case InstBranch(dest)                                        => tls ??= dest
    case InstBranch2(cond, ifTrue, ifFalse)                      => tls ?= cond ??= ifTrue ??= ifFalse
    case i @ InstSwitch(opndTy, opnd, defDest, cases)            => tls += opndTy ?= opnd ??= defDest ??= cases.map(_._1)
    case InstCall(sig, callee, argList, excClause, keepalives)   => tls += sig ?= callee ??= argList ??= excClause
    case InstTailCall(sig, callee, argList)                      => tls += sig ?= callee ??= argList
    case InstRet(funcVer, retVals)                               => tls ??= retVals
    case InstThrow(excVal)                                       => tls ?= excVal
    case InstExtractValue(strTy, index, opnd)                    => tls += strTy ?= opnd
    case InstInsertValue(strTy, index, opnd, newVal)             => tls += strTy ?= opnd ?= newVal
    case InstExtractElement(seqTy, indTy, opnd, index)           => tls += seqTy += indTy ?= opnd ?= index
    case InstInsertElement(seqTy, indTy, opnd, index, newVal)    => tls += seqTy += indTy ?= opnd ?= index ?= newVal
    case InstShuffleVector(vecTy, maskTy, vec1, vec2, mask)      => tls += vecTy += maskTy ?= vec1 ?= vec2 ?= mask
    case InstNew(allocTy, excClause)                             => tls += allocTy ??= excClause
    case InstNewHybrid(allocTy, lenTy, length, excClause)        => tls += allocTy += lenTy ?= length ??= excClause
    case InstAlloca(allocTy, excClause)                          => tls += allocTy ??= excClause
    case InstAllocaHybrid(allocTy, lenTy, length, excClause)     => tls += allocTy += lenTy ?= length ??= excClause
    case InstGetIRef(referentTy, opnd)                           => tls += referentTy ?= opnd
    case InstGetFieldIRef(ptr, referentTy, index, opnd)          => tls += referentTy ?= opnd
    case InstGetElemIRef(ptr, referentTy, indTy, opnd, index)    => tls += referentTy += indTy ?= opnd ?= index
    case InstShiftIRef(ptr, referentTy, offTy, opnd, offset)     => tls += referentTy += offTy ?= opnd ?= offset
    case InstGetVarPartIRef(ptr, referentTy, opnd)               => tls += referentTy ?= opnd
    case InstLoad(ptr, ord, referentTy, loc, excClause)          => tls += referentTy ?= loc ??= excClause
    case InstStore(ptr, ord, referentTy, loc, newVal, excClause) => tls += referentTy ?= loc ?= newVal ??= excClause
    case InstCmpXchg(ptr, weak, ordSucc, ordFail, referentTy, loc, expected, desired, excClause) // This line is long ...
    => tls += referentTy ?= loc ?= expected ?= desired ??= excClause
    case InstAtomicRMW(ptr, ord, op, referentTy, loc, opnd, excClause) => tls += referentTy ?= loc ?= opnd ??= excClause
    case InstFence(ord) =>
    case InstTrap(retTys, excClause, keepalives) => tls ++= retTys ??= excClause
    case InstWatchPoint(wpID, retTys, dis, ena, exc, keepalives) => tls ++= retTys ??= dis ??= ena ????= exc
    case InstWPBranch(wpID, dis, ena) => tls ??= dis ??= ena
    case InstCCall(callConv, funcTy, sig, callee, argList, excClause, keepalives) // This line is long, too...
    => tls += funcTy += sig ?= callee ??= argList ??= excClause
    case InstNewThread(stack, threadLocal, newStackAction, excClause) => {
      tls ?= stack ???= threadLocal ??= excClause
      visitNewStackAction(newStackAction)
    }
    case InstSwapStack(swappee, curStackAction, newStackAction, excClause, keepalives) => {
      tls ?= swappee ??= excClause
      curStackAction match {
        case RetWith(ts) => tls ++= ts
        case KillOld()   =>
      }
      visitNewStackAction(newStackAction)
    }
    case InstCommInst(ci, flagList, typeList, funcSigList, argList, excClause, keepalives) =>
      tls ++= typeList ++= funcSigList ??= argList ??= excClause

    case _ => throw new BootImageBuilderException("Unknown instruction: " + inst.getClass.getName)
  }

  private def visitNewStackAction(nsa: NewStackAction): Unit = nsa match {
    case PassValues(ts, vs) => tls ++= ts ??= vs
    case ThrowExc(e)        =>
  }

  /**
   * Visit an allocation unit (global cell or heap object). As a pre-condition, begin must refer to the beginning of
   * the unit.
   *
   * @param begin The beginning of the object.
   */
  private def visitAllocUnit(begin: Word): Unit = {
    logger.debug("Visiting allocation unit 0x%x".format(begin))
    if (globalMemory.isInSpace(begin)) {
      val GlobalCellRec(begin2, end, g) = getGlobalCellRecExact(begin)
      assert(begin2 == begin)
      logger.debug("  It's global cell. begin=0x%x, end=0x%x, type=%s".format(begin, end, g.cellTy))
      tls += g
      visitAllocUnitFields(begin, g.cellTy)
    } else if (smallObjectSpace.isInSpace(begin) || largeObjectSpace.isInSpace(begin)) {
      val objRef = begin
      val tag = HeaderUtils.getTag(objRef)
      val ty = HeaderUtils.getType(microVM, tag)
      tls += ty
      visitAllocUnitFields(objRef, ty)
    } else {
      throw new BootImageBuilderException("Address %d 0x%x is not in the heap or global space".format(begin, begin))
    }
  }

  private def visitAllocUnitFields(begin: Word, ty: Type): Unit = {
    MemoryDataScanner.scanAllocUnit(begin, this)
  }

  override def visitRefField(objRef: Word, iRef: Word, toObj: Word, isWeak: Boolean): Unit = {
    if (toObj != 0L) {
      allocs += toObj
    }
  }

  override def visitIRefField(objRef: Word, iRef: Word, toObj: Word, toOffset: Word): Unit = {
    if (toObj == 0L) {
      if (globalMemory.isInSpace(toOffset)) {
        val gr = getGlobalCellRec(toOffset)
        tls += gr.g
      } else {
        if (largeObjectSpace.isInSpace(toOffset)) {
          throw new BootImageBuilderException(("Error: non-object iref to LOS found. This is probably a reference to " +
            "a stack cell (ALLOCA). It should not exist during boot image building. source: [0x%x + 0x%x] -> dest 0x%x").format(
              objRef, iRef, toOffset))
        } else {
          throw new BootImageBuilderException(("Error: non-object iref not referring to global. source: [0x%x + 0x%x] -> dest 0x%x").format(
            objRef, iRef, toOffset))
        }
      }
    } else {
      if (toObj != 0L) {
        allocs += toObj
      }
    }
  }

  override def visitTagRefField(objRef: Word, iRef: Word, toObj: Word): Unit = {
    if (toObj != 0L) {
      allocs += toObj
    }
  }

  override def visitFuncRefField(objRef: Word, iRef: Word, toFunc: Option[Function]): Unit = {
    toFunc.foreach(tls+=)
  }

  override def visitUPtrField(objRef: Word, iRef: Word, toAddr: Word): Unit = {
    maybeGetGlobalCellRec(toAddr).foreach{ gcr => 
      tls += gcr.g
    }
  }
  override def visitUFPField(objRef: Word, iRef: Word, toAddr: Word): Unit = {}
  override def visitStackRefField(objRef: Word, iRef: Word, toStack: Option[InterpreterStack]): Unit =
    toStack.foreach(o => throw noReach("stack", objRef, iRef, o))
  override def visitThreadRefField(objRef: Word, iRef: Word, toThread: Option[InterpreterThread]): Unit =
    toThread.foreach(o => throw noReach("thread", objRef, iRef, o))
  override def visitFCRefField(objRef: Word, iRef: Word, toFCRef: Option[FrameCursor]): Unit =
    toFCRef.foreach(o => throw noReach("frame cursor", objRef, iRef, o))
  override def visitIRBuilderRefField(objRef: Word, iRef: Word, toIRNode: Option[IRBuilder]): Unit =
    toIRNode.foreach(o => throw noReach("IR builder", objRef, iRef, o))

  private def noReach(kind: String, objRef: Word, iRef: Word, obj: AnyRef) =
    new BootImageBuilderException(
      "Error: Non-null %s field is not allowed in boot image. From: obj 0x%x, iref 0x%x. To stack id: %s".format(
        objRef, iRef, obj))
}