# 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
{.push checks:off.}   # The VM is a critical point where checks are deleterious

import std/monotimes
import std/math


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


import strutils
when debugVM:
    import std/strformat
    import std/terminal



type
    PeonVM* = ref object
        ## The Peon Virtual Machine.
        ## Note how the only data
        ## type we handle here is
        ## a 64-bit unsigned integer:
        ## This is to allow the use
        ## of unboxed primitive types.
        ## For more complex types, the
        ## value represents a pointer to
        ## some stack- or heap-allocated
        ## object. The VM has no concept
        ## of type by itself and it relies
        ## on the compiler to produce the
        ## correct results
        ip: uint64                  # Instruction pointer
        chunk: Chunk                # Piece of bytecode to execute
        calls: seq[uint64]          # Our call stack
        operands: seq[uint64]       # Our operand stack
        cache: array[6, uint64]     # Singletons cache
        frames: seq[uint64]         # Stores the bottom of stack frames
        closedOver: seq[uint64]     # Stores variables that do not have stack semantics
        results: seq[uint64]        # Stores function's results (return values)


proc initCache*(self: PeonVM) =
    ## Initializes the VM's
    ## singletons cache
    self.cache[0] = 0x0  # Nil
    self.cache[1] = 0x1  # True
    self.cache[2] = 0x2  # False
    self.cache[3] = 0x3  # Positive inf
    self.cache[4] = 0x4  # Negative inf
    self.cache[5] = 0x5  # NaN


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


# Getters for singleton types
{.push inline.}

proc getNil*(self: PeonVM): uint64 = self.cache[2]

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

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

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


# Thanks to nim's *genius* idea of making x !> y a template
# for y < x (which by itself is fine) together with the fact
# that the order of evaluation of templates with the same
# expression is fucking stupid (see https://nim-lang.org/docs/manual.html#order-of-evaluation
# and https://github.com/nim-lang/Nim/issues/10425 and try not to 
# bang your head against the nearest wall), we need a custom operator
# that preserves the natural order of evaluation
proc `!>`[T](a, b: T): auto {.inline.} = 
    b < a


proc `!>=`[T](a, b: T): auto {.inline, used.} = 
    b <= a


# Stack primitives. Note: all accesses to the call stack
# that go through the getc/setc 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: uint64)  =
    ## Pushes a value object onto the
    ## operand stack
    self.operands.add(obj)


proc pop(self: PeonVM): uint64  =
    ## Pops a value off the
    ## operand stack and
    ## returns it
    return self.operands.pop()


proc peekb(self: PeonVM, distance: BackwardsIndex = ^1): uint64  =
    ## Returns the value at the
    ## given (backwards) distance from the top of
    ## the operand stack without consuming it
    return self.operands[distance]


proc peek(self: PeonVM, distance: int = 0): uint64  =
    ## Returns the value at the
    ## given distance from the top of
    ## the operand stack without consuming it
    if distance < 0:
        return self.peekb(^(-distance))
    return self.operands[self.operands.high() + distance]



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


proc popc(self: PeonVM): uint64  =
    ## Pops a value off the call
    ## stack and returns it
    return self.calls.pop()


proc peekc(self: PeonVM, distance: int = 0): uint64 {.used.} =
    ## Returns the value 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: uint): uint64  =
    ## Accessor method that abstracts
    ## indexing our call stack through stack
    ## frames
    return self.calls[idx + self.frames[^1]]


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

# Byte-level primitives to read and 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())


# Functions to read primitives from the chunk's
# constants table

proc constReadInt64(self: PeonVM, idx: int): int64  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as 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],
            ]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadUInt64(self: PeonVM, idx: int): uint64  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as 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],
            ]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadUInt32(self: PeonVM, idx: int): uint32  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an int32
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3]]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadInt32(self: PeonVM, idx: int): int32  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an uint32
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3]]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadInt16(self: PeonVM, idx: int): int16  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an int16
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1]]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadUInt16(self: PeonVM, idx: int): uint16  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an uint16
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1]]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadInt8(self: PeonVM, idx: int): int8  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an int8
    result = int8(self.chunk.consts[idx])


proc constReadUInt8(self: PeonVM, idx: int): uint8  =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as an uint8
    result = self.chunk.consts[idx]


proc constReadFloat32(self: PeonVM, idx: int): float32 =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as a float32
    var arr = [self.chunk.consts[idx], self.chunk.consts[idx + 1],
               self.chunk.consts[idx + 2], self.chunk.consts[idx + 3]]
    copyMem(result.addr, arr.addr, sizeof(arr))


proc constReadFloat64(self: PeonVM, idx: int): float =
    ## Reads a constant from the
    ## chunk's constant table and
    ## returns it as a float
    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]]
    copyMem(result.addr, arr.addr, sizeof(arr))
{.pop.}


when debugVM:   # So nim shuts up 
    proc debug(self: PeonVM) =
        ## Implements the VM's runtime
        ## debugger
        styledEcho fgMagenta, "IP: ", fgYellow, &"{self.ip}"
        styledEcho fgBlue, "Instruction: ", fgRed, &"{OpCode(self.chunk.code[self.ip])} (", fgYellow, $self.chunk.code[self.ip], fgRed, ")"
        if self.calls.len() !> 0:
            stdout.styledWrite(fgGreen, "Call Stack: ", fgMagenta, "[")
            for i, e in self.calls:
                stdout.styledWrite(fgYellow, $e)
                if i < self.calls.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
        if self.operands.len() !> 0:
            stdout.styledWrite(fgBlue, "Operand Stack: ", fgMagenta, "[")
            for i, e in self.operands:
                stdout.styledWrite(fgYellow, $e)
                if i < self.operands.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
        if self.frames.len() !> 0:
            stdout.styledWrite(fgCyan, "Current Frame: ", fgMagenta, "[")
            for i, e in self.calls[self.frames[^1]..^1]:
                stdout.styledWrite(fgYellow, $e)
                if i < self.calls.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
            stdout.styledWrite(fgRed, "Live stack frames: ", fgMagenta, "[")
            for i, e in self.frames:
                stdout.styledWrite(fgYellow, $e)
                if i < self.frames.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
        if self.closedOver.len() !> 0:
            stdout.styledWrite(fgGreen, "Closure Array: ", fgMagenta, "[")
            for i, e in self.closedOver:
                stdout.styledWrite(fgYellow, $e)
                if i < self.closedOver.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
        if self.results.len() !> 0:
            stdout.styledWrite(fgYellow, "Function Results: ", fgMagenta, "[")
            for i, e in self.results:
                stdout.styledWrite(fgYellow, $e)
                if i < self.results.high():
                    stdout.styledWrite(fgYellow, ", ")
            styledEcho fgMagenta, "]"
        discard readLine stdin


proc dispatch*(self: PeonVM) =
    ## Main bytecode dispatch loop
    var instruction {.register.}: OpCode
    while true:
        {.computedgoto.}  # https://nim-lang.org/docs/manual.html#pragmas-computedgoto-pragma
        when debugVM:
            self.debug()
        instruction = OpCode(self.readByte())
        case instruction:
            # Constant loading instructions
            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(uint64(self.constReadInt64(int(self.readLong()))))
            of LoadUInt64:
                self.push(uint64(self.constReadUInt64(int(self.readLong()))))
            of LoadUInt32:
                self.push(uint64(self.constReadUInt32(int(self.readLong()))))
            of LoadInt32:
                self.push(uint64(self.constReadInt32(int(self.readLong()))))
            of LoadInt16:
                self.push(uint64(self.constReadInt16(int(self.readLong()))))
            of LoadUInt16:
                self.push(uint64(self.constReadUInt16(int(self.readLong()))))
            of LoadInt8:
                self.push(uint64(self.constReadInt8(int(self.readLong()))))
            of LoadUInt8:
                self.push(uint64(self.constReadUInt8(int(self.readLong()))))
            of LoadString:
                # TODO: Use constReadString with own memory manager
                # Strings are broken rn!!
                let size = int(self.readLong())
                let idx = int(self.readLong())
                var str = self.chunk.consts[idx..<idx + size].fromBytes()
                self.push(cast[uint64](str.addr))
            # We cast instead of converting because, unlike with integers,
            # we don't want nim to touch any of the bits of the underlying
            # value!
            of LoadFloat32:
                self.push(cast[uint64](self.constReadFloat32(int(self.readLong()))))
            of LoadFloat64:
                self.push(cast[uint64](self.constReadFloat64(int(self.readLong()))))
            of LoadFunction:
                # Loads a function address onto the operand stack
                self.push(uint64(self.readLong()))
            of LoadReturnAddress:
                # Loads a 32-bit unsigned integer onto the operand stack.
                # Used to load function return addresses
                self.push(uint64(self.readUInt()))
            of Call:
                # Calls a peon function. The calling convention here
                # is pretty simple: the first value in the frame is
                # the new instruction pointer to jump to, then a
                # 32-bit return address follows. After that, all
                # arguments and locals follow. Note that, due to
                # how the stack works, all arguments before the call 
                # are in the reverse order in which they are passed
                # to the function
                let argc = self.readLong().int
                let retAddr = self.peek(-argc - 1)    # Return address
                let jmpAddr = self.peek(-argc - 2)    # Function address
                self.ip = jmpAddr
                self.pushc(jmpAddr)
                self.pushc(retAddr)
                # Creates a new result slot for the
                # function's return value
                self.results.add(self.getNil())
                # Creates a new call frame
                self.frames.add(uint64(self.calls.len() - 2))
                # Loads the arguments onto the stack
                for _ in 0..<argc:
                    self.pushc(self.pop())
                # Pops the function and return address 
                # off the operand stack since they're 
                # not needed there anymore
                discard self.pop()
                discard self.pop()
            of 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
                while self.calls.len().uint64 !> self.frames[^1] + 2'u64:
                    # Discards the function's local variables,
                    # if there is any
                    discard self.popc()
                let ret = self.popc() # Return address
                discard self.popc()   # Function address
                if self.readByte() == 1:
                    # Function is non-void!
                    self.push(self.results.pop())
                else:
                    discard self.results.pop()
                # Discard the topmost stack frame
                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. A Return
                # instruction will pop this
                # off the results array and
                # onto the operand stack when
                # the current function exits.
                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 = self.readLong()
                when debugVM:
                    assert idx.int - self.calls.high() <= 1, "StoreVar index is bigger than the length of the call stack"
                if idx + self.frames[^1] <= self.calls.high().uint:
                    self.setc(idx, self.pop())
                else:
                    self.pushc(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 LoadVar:
                # Pushes a variable onto the operand
                # stack
                self.push(self.getc(self.readLong()))
            of NoOp:
                # Does nothing
                continue
            of PopC:
                # Pops a value off the call stack
                discard self.popc()
            of Pop:
                # Pops a value off the operand stack
                discard self.pop()
            of PushC:
                # Pushes a value from the operand stack
                # onto the call stack
                self.pushc(self.pop())
            of PopRepl:
                # Pops a peon object off the
                # operand stack and prints it.
                # Used in interactive REPL mode
                if self.frames.len() !> 1:
                    discard self.pop()
                    continue
                echo self.pop()
            of PopN:
                # Pops N elements off the call stack
                for _ in 0..<int(self.readShort()):
                    discard self.popc()
            # Jump opcodes
            of Jump:
                # Absolute jump
                self.ip = self.readLong()
            of JumpForwards:
                # Relative, forward-jump
                self.ip += self.readLong()
            of JumpBackwards:
                # Relative, backward-jump
                self.ip -= self.readLong()
            of JumpIfFalse:
                # Conditional positive jump
                if not self.peek().bool:
                    self.ip += self.readLong()
            of JumpIfTrue:
                # Conditional positive jump
                if self.peek().bool:
                    self.ip += self.readLong()
            of JumpIfFalsePop:
                let ip = self.readLong()
                if not self.peek().bool:
                    self.ip += ip
                discard self.pop()
            of JumpIfFalseOrPop:
                if not self.peek().bool:
                    self.ip += self.readLong()
                else:
                    discard self.pop()
            # Built-in operations on primitive types.
            # Note: for operations where the order of
            # the operands matters, we don't need to
            # swap the order of the calls to pop: this 
            # is because operators are handled like peon 
            # functions, which means the arguments are 
            # already reversed on the stack when we 
            # execute the instruction
            of Negate:
                self.push(uint64(-int64(self.pop())))
            of NegateFloat64:
                self.push(cast[uint64](-cast[float](self.pop())))
            of NegateFloat32:
                self.push(cast[uint64](-cast[float32](self.pop())))
            of Add:
                self.push(self.pop() + self.pop())
            of Subtract:
                self.push(self.pop() - self.pop())
            of Multiply:
                self.push(self.pop() * self.pop())
            of Divide:
                self.push(self.pop() div self.pop())
            of SignedDivide:
                self.push(uint64(int64(self.pop()) div int64(self.pop())))
            of AddFloat64:
                self.push(cast[uint64](cast[float](self.pop()) + cast[float](self.pop())))
            of SubtractFloat64:
                self.push(cast[uint64](cast[float](self.pop()) - cast[float](self.pop())))
            of MultiplyFloat64:
                self.push(cast[uint64](cast[float](self.pop()) * cast[float](self.pop())))
            of DivideFloat64:
                self.push(cast[uint64](cast[float](self.pop()) / cast[float](self.pop())))
            of AddFloat32:
                self.push(cast[uint64](cast[float32](self.pop()) + cast[float32](self.pop())))
            of SubtractFloat32:
                self.push(cast[uint64](cast[float32](self.pop()) - cast[float32](self.pop())))
            of MultiplyFloat32:
                self.push(cast[uint64](cast[float32](self.pop()) * cast[float32](self.pop())))
            of DivideFloat32:
                self.push(cast[uint64](cast[float32](self.pop()) / cast[float32](self.pop())))
            of Pow:
                self.push(uint64(self.pop() ^ self.pop()))
            of SignedPow:
                self.push(uint64(int64(self.pop()) ^ int64(self.pop())))
            of PowFloat64:
                self.push(cast[uint64](pow(cast[float](self.pop()), cast[float](self.pop()))))
            of PowFloat32:
                self.push(cast[uint64](pow(cast[float](self.pop()), cast[float](self.pop()))))
            of Mod:
                self.push(uint64(self.pop() mod self.pop()))
            of SignedMod:
                self.push(uint64(int64(self.pop()) mod int64(self.pop())))
            of ModFloat64:
                self.push(cast[uint64](floorMod(cast[float](self.pop()), cast[float](self.pop()))))
            of ModFloat32:
                self.push(cast[uint64](floorMod(cast[float](self.pop()), cast[float](self.pop()))))
            of LShift:
                self.push(self.pop() shl self.pop())
            of RShift:
                self.push(self.pop() shr self.pop())
            of Xor:
                self.push(self.pop() xor self.pop())
            of Not:
                self.push(not self.pop())
            of And:
                self.push(self.pop() and self.pop())
            # Comparison opcodes
            of Equal:
                self.push(self.getBool(self.pop() == self.pop()))
            of NotEqual:
                self.push(self.getBool(self.pop() != self.pop()))
            of GreaterThan:
                self.push(self.getBool(self.pop() !> self.pop()))
            of LessThan:
                self.push(self.getBool(self.pop() < self.pop()))
            of GreaterOrEqual:
                self.push(self.getBool(self.pop() !>= self.pop()))
            of LessOrEqual:
                self.push(self.getBool(self.pop() <= self.pop()))
            # Print opcodes
            of PrintInt64:
                echo int64(self.pop())
            of PrintUInt64:
                echo self.pop()
            of PrintInt32:
                echo int32(self.pop())
            of PrintUInt32:
                echo uint32(self.pop())
            of PrintInt16:
                echo int16(self.pop())
            of PrintUInt16:
                echo uint16(self.pop())
            of PrintInt8:
                echo int8(self.pop())
            of PrintUInt8:
                echo uint8(self.pop())
            of PrintFloat32:
                echo cast[float32](self.pop())
            of PrintFloat64:
                echo cast[float](self.pop())
            of PrintHex:
                echo "0x" & self.pop().toHex().strip(chars={'0'})
            of PrintBool:
                if self.pop().bool:
                    echo "true"
                else:
                    echo "false"
            of PrintInf:
                if self.pop() == 0x3:
                    echo "-inf"
                else:
                    echo "inf"
            of PrintNan:
                echo "nan"
            of PrintString:
                echo cast[ptr string](self.pop())[]  # TODO
            of SysClock64:
                # Pushes the value of a monotonic clock
                # onto the operand stack. This can be used
                # to track system time accurately, but it 
                # cannot be converted to a date. The number
                # is in seconds
                self.push(cast[uint64](getMonoTime().ticks().float() / 1_000_000_000))
            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()

{.pop.}