# 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 strutils
import strformat

import types
import ../config
import ../frontend/meta/bytecode
import ../util/multibyte

export types

type
    PeonVM* = ref object
        ## The Peon Virtual Machine
        calls: seq[PeonObject]      # Our call stack
        operands: seq[PeonObject]   # Our operand stack
        ip: uint32                  # Instruction pointer
        cache: array[6, PeonObject] # Singletons cache
        chunk: Chunk                # Piece of bytecode to execute
        frames: seq[int]            # Stores the bottom of stack frames
        closedOver: seq[PeonObject] # Stores variables that do not have stack semantics
        results: seq[PeonObject]    # Stores function's results


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.calls = newSeq[PeonObject]()
    result.operands = newSeq[PeonObject]()
    result.initCache()


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

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

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

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

proc getNan*(self: PeonVM): PeonObject {.inline.} = 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
    ## operand stack
    self.operands.add(obj)


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


proc peek(self: PeonVM, distance: int = 0): PeonObject =
    ## Returns the Peon object at the
    ## given distance from the top of
    ## the operand stack without consuming it
    return self.operands[self.operands.high() + distance]


proc pushc(self: PeonVM, val: PeonObject) =
    ## Pushes a new object to the
    ## call stack
    self.calls.add(val)


proc popc(self: PeonVM): PeonObject =
    ## Pops an object off the call
    ## stack and returns it
    return self.calls.pop()


proc peekc(self: PeonVM, distance: int = 0): PeonObject =
    ## Returns the Peon object at the
    ## given distance from the top of
    ## the call stack without consuming it
    return self.calls[self.calls.high() + distance]


proc getc(self: PeonVM, idx: int): PeonObject =
    ## Accessor method that abstracts
    ## indexing our call stack through stack
    ## frames
    return self.calls[idx + self.frames[^1]]


proc setc(self: PeonVM, idx: int, val: PeonObject) =
    ## Setter method that abstracts
    ## indexing our call stack through stack
    ## frames
    self.calls[idx + self.frames[^1]] = val

## Byte-level primitives to read/decode
## bytecode

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 readUInt(self: PeonVM): uint32 =
    ## Reads three bytes from the
    ## bytecode and returns them
    ## as an unsigned 32 bit
    ## integer
    return uint32([self.readByte(), self.readByte(), self.readByte(), self.readByte()].fromQuad())


proc constReadInt64(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 constReadUInt64(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 constReadUInt32(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 constReadInt32(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 constReadInt16(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an Int16
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1]]
    result = PeonObject(kind: Int16)
    copyMem(result.short.addr, arr.addr, sizeof(arr))


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


proc constReadInt8(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an Int8
    result = PeonObject(kind: Int8, tiny: int8(self.chunk.consts[idx]))


proc constReadUInt8(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is an UInt8
    result = PeonObject(kind: UInt8, uTiny: self.chunk.consts[idx])


proc constReadString(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is a string
    let size = self.readLong()
    result = PeonObject(kind: String, str: self.chunk.consts[idx..<size].fromBytes())


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


proc constReadFloat64(self: PeonVM, idx: int): PeonObject =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns a Peon object. Assumes
    ## the constant is a Float64
    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: Float64)
    copyMem(result.`float`.addr, arr.addr, sizeof(arr))


proc dispatch*(self: PeonVM) =
    ## Main bytecode dispatch loop
    var instruction {.register.}: OpCode
    while true:
        when DEBUG_TRACE_VM:
            echo &"IP: {self.ip}"
            echo &"Instruction: {OpCode(self.chunk.code[self.ip])}"
            if self.calls.len() > 0:
                echo &"Call Stack: {self.calls}"
            if self.operands.len() > 0:
                echo &"Operand Stack: {self.operands}"
            if self.frames.len() > 0:
                echo &"Current Frame: {self.calls[self.frames[^1]..^1]}"
                echo &"Frames: {self.frames}"
            if self.closedOver.len() > 0:
                echo &"Closure Array: {self.closedOver}"
            if self.results.len() > 0:
                echo &"Results: {self.results}"
            discard readLine stdin
        instruction = OpCode(self.readByte())
        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.constReadInt64(int(self.readLong())))
            of LoadUInt64:
                self.push(self.constReadUInt64(int(self.readLong())))
            of LoadUInt32:
                self.push(self.constReadUInt32(int(self.readLong())))
            of LoadInt32:
                self.push(self.constReadInt32(int(self.readLong())))
            of LoadInt16:
                self.push(self.constReadInt16(int(self.readLong())))
            of LoadUInt16:
                self.push(self.constReadUInt16(int(self.readLong())))
            of LoadInt8:
                self.push(self.constReadInt8(int(self.readLong())))
            of LoadUInt8:
                self.push(self.constReadUInt8(int(self.readLong())))
            of LoadString:
                self.push(self.constReadString(int(self.readLong())))
            of LoadFloat32:
                self.push(self.constReadFloat32(int(self.readLong())))
            of LoadFloat64:
                self.push(self.constReadFloat64(int(self.readLong())))
            of LoadFunction:
                self.pushc(PeonObject(kind: Function, ip: self.readLong()))
            of LoadFunctionObj:
                self.push(PeonObject(kind: Function, ip: self.readLong()))
            of LoadReturnAddress:
                self.pushc(PeonObject(kind: UInt32, uInt: self.readUInt()))
            of Call:
                # Calls a function. The calling convention for peon
                # functions is pretty simple: the first item in the
                # frame is a function object which contains the new
                # instruction pointer to jump to, followed by a 32-bit
                # return address. After that, all arguments and locals 
                # follow
                var size {.used.} = self.readLong().int
                self.frames.add(self.calls.len() - 2)
                self.ip = self.peekc(-1).ip
                self.results.add(self.getNil())
                # TODO: Use the frame size once
                # we have more control over the
                # memory
                #[while size > 0:
                    dec(size)
                    self.pushc(self.getNil())
                ]#
            of LoadArgument:
                self.pushc(self.pop())
            of OpCode.Return:
                # Returns from a function.
                # Every peon program is wrapped
                # in a hidden function, so this 
                # will also exit the VM if we're 
                # at  the end of the program
                let ret = self.popc()
                discard self.popc()   # Function object
                if self.readByte() == 1:
                    # Function is non-void!
                    self.push(self.results.pop())
                else:
                    discard self.results.pop()
                discard self.frames.pop()
                if self.frames.len() == 0:
                    # End of the program!
                    return
                self.ip = ret.uInt
            of SetResult:
                # Sets the result of the
                # current function
                self.results[self.frames.high()] = self.pop()
            of StoreVar:
                # Stores the value at the top of the operand stack
                # into the given call stack index
                let idx = int(self.readLong()) + self.frames[^1]
                if idx <= self.calls.high():
                    self.setc(idx, self.pop())
                else:
                    self.calls.add(self.pop())
            of StoreClosure:
                # Stores/updates the value of a closed-over
                # variable
                let idx = self.readLong().int
                if idx > self.closedOver.high():
                    # Note: we *peek* the stack, but we
                    # don't pop!
                    self.closedOver.add(self.peek())
                else:
                    self.closedOver[idx] = self.peek()
            of LoadClosure:
                # Loads a closed-over variable onto the
                # stack
                self.push(self.closedOver[self.readLong()])
            of PopClosure:
                # Removes a closed-over variable from the closure
                # array
                discard self.closedOver.pop()
            of LoadVar:
                self.push(self.getc(int(self.readLong())))
            of NoOp:
                continue
            of PopC:
                discard self.popc()
            of Pop:
                discard self.pop()
            of PushC:
                self.pushc(self.pop())
            of PopRepl:
                if self.frames.len() > 1:
                    discard self.pop()
                    continue
                let popped = self.pop()
                case popped.kind:
                of Int64:
                    echo &"{popped.long}'i64"
                of UInt64:
                    echo &"{popped.uLong}'u64"
                of Int32:
                    echo &"{popped.`int`}'i32"
                of UInt32:
                    echo &"{popped.uInt}'u32"
                of Int16:
                    echo &"{popped.short}'i16"
                of UInt16:
                    echo &"{popped.uShort}'u16"
                of Int8:
                    echo &"{popped.tiny}'i8"
                of UInt8:
                    echo &"{popped.uTiny}'u8"
                of Float32:
                    echo &"{popped.halfFloat}'f32"
                of Float64:
                    echo &"{popped.`float`}'f64"
                of ObjectKind.Inf:
                    if popped.positive:
                        echo "inf"
                    else:
                        echo "-inf"
                of ObjectKind.Nan, Nil:
                    echo ($popped.kind).toLowerAscii()
                else:
                    discard
            of PopN:
                for _ in 0..<int(self.readShort()):
                    discard self.popc()
            # Jump opcodes
            of Jump:
                self.ip = self.readLong()
            of JumpForwards:
                self.ip += self.readLong()
            of JumpBackwards:
                self.ip -= self.readLong()
            of JumpIfFalse:
                if not self.peek().boolean:
                    self.ip += self.readLong()
            of JumpIfTrue:
                if self.peek().boolean:
                    self.ip += self.readLong()
            of JumpIfFalsePop:
                let ip = self.readLong()
                if not self.peek().boolean:
                    self.ip += ip
                discard self.pop()
            of JumpIfFalseOrPop:
                if not self.peek().boolean:
                    self.ip += self.readLong()
                else:
                    discard self.pop()
            # Built-in operations on primitive types
            of AddInt64:
                self.push(PeonObject(kind: Int64, long: self.pop().long + self.pop().long))
            of SubInt64:
                let second = self.pop()
                self.push(PeonObject(kind: Int64, long: self.pop().long - second.long))
            of MulInt64:
                self.push(PeonObject(kind: Int64, long: self.pop().long * self.pop().long))
            of DivInt64:
                let second = self.pop()
                self.push(PeonObject(kind: Int64, long: self.pop().long div second.long))
            of AddUInt64:
                self.push(PeonObject(kind: UInt64, uLong: self.pop().uLong + self.pop().uLong))
            of SubUInt64:
                let second = self.pop()
                self.push(PeonObject(kind: UInt64, uLong: self.pop().uLong - second.uLong))
            of MulUInt64:
                self.push(PeonObject(kind: UInt64, uLong: self.pop().uLong * self.pop().uLong))
            of DivUInt64:
                let second = self.pop()
                self.push(PeonObject(kind: UInt64, uLong: self.pop().uLong div second.uLong))
            of AddInt32:
                self.push(PeonObject(kind: Int32, `int`: self.pop().`int` + self.pop().`int`))
            of SubInt32:
                let second = self.pop()
                self.push(PeonObject(kind: Int32, `int`: self.pop().`int` - second.`int`))
            of MulInt32:
                self.push(PeonObject(kind: Int32, `int`: self.pop().`int` * self.pop().`int`))
            of DivInt32:
                let second = self.pop()
                self.push(PeonObject(kind: Int32, `int`: self.pop().`int` div second.`int`))
            of AddUInt32:
                self.push(PeonObject(kind: UInt32, uInt: self.pop().uInt + self.pop().uInt))
            of SubUInt32:
                let second = self.pop()
                self.push(PeonObject(kind: UInt32, uInt: self.pop().uInt - second.uInt))
            of MulUInt32:
                self.push(PeonObject(kind: UInt32, uInt: self.pop().uInt * self.pop().uInt))
            of DivUInt32:
                let second = self.pop()
                self.push(PeonObject(kind: UInt32, uInt: self.pop().uInt div second.uInt))
            of AddInt16:
                self.push(PeonObject(kind: Int16, short: self.pop().short + self.pop().short))
            of SubInt16:
                let second = self.pop()
                self.push(PeonObject(kind: Int16, short: self.pop().short - second.short))
            of MulInt16:
                self.push(PeonObject(kind: Int16, short: self.pop().short * self.pop().short))
            of DivInt16:
                let second = self.pop()
                self.push(PeonObject(kind: Int16, short: self.pop().short div second.short))
            of AddUInt16:
                self.push(PeonObject(kind: UInt16, uShort: self.pop().uShort + self.pop().uShort))
            of SubUInt16:
                let second = self.pop()
                self.push(PeonObject(kind: UInt16, uShort: self.pop().uShort - second.uShort))
            of MulUInt16:
                self.push(PeonObject(kind: UInt16, uShort: self.pop().uShort * self.pop().uShort))
            of DivUInt16:
                let second = self.pop()
                self.push(PeonObject(kind: UInt16, uShort: self.pop().uShort div second.uShort))
            of AddInt8:
                self.push(PeonObject(kind: Int8, tiny: self.pop().tiny + self.pop().tiny))
            of SubInt8:
                let second = self.pop()
                self.push(PeonObject(kind: Int8, tiny: self.pop().tiny - second.tiny))
            of MulInt8:
                self.push(PeonObject(kind: Int8, tiny: self.pop().tiny * self.pop().tiny))
            of DivInt8:
                let second = self.pop()
                self.push(PeonObject(kind: Int8, tiny: self.pop().tiny div second.tiny))
            of AddUInt8:
                self.push(PeonObject(kind: UInt8, uTiny: self.pop().uTiny + self.pop().uTiny))
            of SubUInt8:
                let second = self.pop()
                self.push(PeonObject(kind: UInt8, uTiny: self.pop().uTiny - second.uTiny))
            of MulUInt8:
                self.push(PeonObject(kind: UInt8, uTiny: self.pop().uTiny * self.pop().uTiny))
            of DivUInt8:
                let second = self.pop()
                self.push(PeonObject(kind: UInt8, uTiny: self.pop().uTiny div second.uTiny))
            of AddFloat64:
                self.push(PeonObject(kind: Float64, `float`: self.pop().`float` + self.pop().`float`))
            of SubFloat64:
                let second = self.pop()
                self.push(PeonObject(kind: Float64, `float`: self.pop().`float` - second.`float`))
            of MulFloat64:
                self.push(PeonObject(kind: Float64, `float`: self.pop().`float` * self.pop().`float`))
            of DivFloat64:
                let second = self.pop()
                self.push(PeonObject(kind: Float64, `float`: self.pop().`float` / second.`float`))
            of AddFloat32:
                self.push(PeonObject(kind: Float32, halfFloat: self.pop().halfFloat + self.pop().halfFloat))
            of SubFloat32:
                let second = self.pop()
                self.push(PeonObject(kind: Float32, halfFloat: self.pop().halfFloat - second.halfFloat))
            of MulFloat32:
                self.push(PeonObject(kind: Float32, halfFloat: self.pop().halfFloat * self.pop().halfFloat))
            of DivFloat32:
                let second = self.pop()
                self.push(PeonObject(kind: Float32, halfFloat: self.pop().halfFloat / second.halfFloat))
            of NegInt64:
                self.push(PeonObject(kind: Int64, long: -self.pop().long))
            of NegInt32:
                self.push(PeonObject(kind: Int32, `int`: -self.pop().`int`))
            of NegInt16:
                self.push(PeonObject(kind: Int16, short: -self.pop().short))
            of NegInt8:
                self.push(PeonObject(kind: Int8, tiny: -self.pop().tiny))
            else:
                discard


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