# Copyright 2022 Mattia Giambirtone & All Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

## The Peon runtime environment
import types
import ../config
when DEBUG_TRACE_VM:
    import strformat
import ../frontend/meta/bytecode
import ../util/multibyte


export types

type
    PeonVM* = ref object
        ## The Peon Virtual Machine
        stack: seq[PeonObject]
        ip: int                     # Instruction pointer
        cache: array[6, PeonObject] # Singletons cache
        chunk: Chunk                # Piece of bytecode to execute
        frames: seq[int]            # Stores the initial index of stack frames
        heapVars: seq[PeonObject]   # Stores variables that do not have stack semantics (i.e. "static")
        lastPop*: PeonObject         # Used in the REPL


proc initCache*(self: PeonVM) =
    ## Initializes the VM's
    ## singletons cache
    self.cache[0] = PeonObject(kind: Nil)
    self.cache[1] = PeonObject(kind: Bool, boolean: true)
    self.cache[2] = PeonObject(kind: Bool, boolean: false)
    self.cache[3] = PeonObject(kind: ObjectKind.Inf, positive: true)
    self.cache[4] = PeonObject(kind: ObjectKind.Inf, positive: false)
    self.cache[5] = PeonObject(kind: ObjectKind.Nan)


proc newPeonVM*: PeonVM =
    ## Initializes a new, blank VM
    ## for executing Peon bytecode
    new(result)
    result.ip = 0
    result.frames = @[]
    result.stack = newSeq[PeonObject]()
    result.initCache()


## Getters for singleton types (they are cached!)

proc getNil*(self: PeonVM): PeonObject = self.cache[0]

proc getBool*(self: PeonVM, value: bool): PeonObject =
    if value:
        return self.cache[1]
    return self.cache[2]

proc getInf*(self: PeonVM, positive: bool): PeonObject =
    if positive:
        return self.cache[3]
    return self.cache[4]

proc getNan*(self: PeonVM): PeonObject = self.cache[5]

## Stack primitives. Note: all stack accessing that goes
## through the get/set wrappers is frame-relative, meaning
## that the index is added to the current stack frame's 
## bottom to obtain an absolute stack index.

proc push(self: PeonVM, obj: PeonObject) =
    ## Pushes a Peon object onto the
    ## stack
    self.stack.add(obj)


proc pop(self: PeonVM): PeonObject =
    ## Pops a Peon object off the
    ## stack, decreasing the stack
    ## pointer. The object is returned
    return self.stack.pop()


proc peek(self: PeonVM): PeonObject =
    ## Returns the Peon object at the top
    ## of the stack without consuming
    ## it
    return self.stack[^1]


proc get(self: PeonVM, idx: int): PeonObject =
    ## Accessor method that abstracts
    ## stack accessing through stack
    ## frames
    return self.stack[idx + self.frames[^1]]


proc set(self: PeonVM, idx: int, val: PeonObject) =
    ## Setter method that abstracts
    ## stack accessing through stack
    ## frames
    self.stack[idx + self.frames[^1]] = val


proc readByte(self: PeonVM): uint8 =
    ## Reads a single byte from the
    ## bytecode and returns it as an
    ## unsigned 8 bit integer
    inc(self.ip)
    return self.chunk.code[self.ip - 1]


proc readShort(self: PeonVM): uint16 =
    ## Reads two bytes from the
    ## bytecode and returns them
    ## as an unsigned 16 bit
    ## integer
    return [self.readByte(), self.readByte()].fromDouble()


proc readLong(self: PeonVM): uint32 =
    ## Reads three bytes from the
    ## bytecode and returns them
    ## as an unsigned 32 bit
    ## integer. Note however that
    ## the boundary is capped at
    ## 24 bits instead of 32
    return uint32([self.readByte(), self.readByte(), self.readByte()].fromTriple())


proc readInt64(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an Int64
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3],
               self.chunk.consts[idx + 4], self.chunk.consts[idx + 5],
               self.chunk.consts[idx + 6], self.chunk.consts[idx + 7],
            ]
    result = PeonObject(kind: Int64)
    copyMem(result.long.addr, arr.addr, sizeof(arr))


proc readUInt64(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an UInt64
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3],
               self.chunk.consts[idx + 4], self.chunk.consts[idx + 5],
               self.chunk.consts[idx + 6], self.chunk.consts[idx + 7],
            ]
    result = PeonObject(kind: UInt64)
    copyMem(result.uLong.addr, arr.addr, sizeof(arr))


proc readUInt32(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an UInt32
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3]]
    result = PeonObject(kind: UInt32)
    copyMem(result.uInt.addr, arr.addr, sizeof(arr))


proc readInt32(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an Int32
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3]]
    result = PeonObject(kind: Int32)
    copyMem(result.`int`.addr, arr.addr, sizeof(arr))


proc dispatch*(self: PeonVM) =
    ## Main bytecode dispatch loop
    var instruction: OpCode
    while true:
        instruction = OpCode(self.readByte())
        when DEBUG_TRACE_VM:
            echo &"IP: {self.ip}"
            echo &"Instruction: {instruction}"
            echo &"Stack: {self.stack}"
            echo &"Current Frame: {self.stack[self.frames[^1]..^1]}"
            discard readLine stdin
        case instruction:
            # Constant loading
            of LoadTrue:
                self.push(self.getBool(true))
            of LoadFalse:
                self.push(self.getBool(false))
            of LoadNan:
                self.push(self.getNan())
            of LoadNil:
                self.push(self.getNil())
            of LoadInf:
                self.push(self.getInf(true))
            of LoadInt64:
                self.push(self.readInt64(int(self.readLong())))
            of LoadUInt64:
                self.push(self.readUInt64(int(self.readLong())))
            of LoadUInt32:
                self.push(self.readUInt32(int(self.readLong())))
            of Call:
                # Calls a function. The calling convention for peon
                # functions is pretty simple: the return address sits
                # at the bottom of the stack frame, then follow the
                # arguments and all temporaries/local variables
                let newIp = self.readLong()
                # We do this because if we immediately changed
                # the instruction pointer, we'd read the wrong
                # value for the argument count. Storing it and
                # changing it later fixes this issue
                self.frames.add(int(self.readLong()))
                self.ip = int(newIp)
            of OpCode.Return:
                # Returns from a void function or terminates the
                # program entirely if we're at the topmost frame
                let frame = self.frames.pop()
                if self.frames.len() > 1:
                    for i in countdown(self.stack.high(), frame):
                        discard self.pop()
                    self.ip = int(self.pop().uInt)
                else:
                    while self.stack.len() > 0:
                        discard self.pop()
                    return
            of ReturnValue:
                # Returns from a function which has a return value,
                # pushing it on the stack
                let retVal = self.pop()
                let frame = self.frames.pop()
                for i in countdown(frame, 1):
                    discard self.pop()
                self.ip = int(self.pop().uInt)
                self.push(retVal)
            of StoreVar:
                # Stores the value at the top of the stack
                # into the given stack index
                self.set(int(self.readLong()), self.pop())
            of StoreHeap:
                # Stores/updates the value of a closed-over
                # variable
                let idx = self.readLong().int
                if idx > self.heapVars.high():
                    self.heapVars.add(self.pop())
                else:
                    self.heapVars[idx] = self.pop()
            of LoadHeap:
                self.push(self.heapVars[self.readLong()])
            of LoadVar:
                self.push(self.get(int(self.readLong())))
            of NoOp:
                continue
            of Pop:
                self.lastPop = self.pop()
            of PopN:
                for _ in 0..<int(self.readShort()):
                    discard self.pop()
            of Jump:
                self.ip = int(self.readShort())
            of JumpForwards:
                self.ip += int(self.readShort())
            of JumpBackwards:
                self.ip -= int(self.readShort())
            of JumpIfFalse:
                if not self.peek().boolean:
                    self.ip += int(self.readShort())
            of JumpIfTrue:
                if self.peek().boolean:
                    self.ip += int(self.readShort())
            of JumpIfFalsePop:
                if not self.peek().boolean:
                    self.ip += int(self.readShort())
                discard self.pop()
            of JumpIfFalseOrPop:
                if not self.peek().boolean:
                    self.ip += int(self.readShort())
                else:
                    discard self.pop()
            of LongJumpIfFalse:
                if not self.peek().boolean:
                    self.ip += int(self.readLong())
            of LongJumpIfFalsePop:
                if not self.peek().boolean:
                    self.ip += int(self.readLong())
                discard self.pop()
            of LongJumpForwards:
                self.ip += int(self.readLong())
            of LongJumpBackwards:
                self.ip -= int(self.readLong())
            of LongJump:
                self.ip = int(self.readLong())
            of LongJumpIfFalseOrPop:
                if not self.peek().boolean:
                    self.ip += int(self.readLong())
                else:
                    discard self.pop()
            else:
                discard


proc run*(self: PeonVM, chunk: Chunk) =
    ## Executes a piece of Peon bytecode.
    self.chunk = chunk
    self.frames = @[0]
    self.stack = @[]
    self.ip = 0
    self.lastPop = self.getNil()
    self.dispatch()