{-
-- EPITECH PROJECT, 2025
-- Glados
-- File description:
-- Function Instructions
-}

{-|
Module      : VM.Instruction.Function
Description : Implementation of function call and return operations.
Stability   : stable

Handles the creation and destruction of stack frames (CallSnapshots).
Ensures proper scope management and control flow between functions.
-}
module VM.Instruction.Function
    ( instCall
    , instTailCall
    , instCallIndirect
    , instRet
    , instMakeClosure
    , instGetFuncAddr
    ) where

import Control.Monad.State.Strict (get, put)
import qualified Data.Vector as V

import VM.VMState (VirtualMachine, VMState(..), doSnapshot)
import VM.VMValue (VMValue(..))
import VM.CallSnapshot (CallSnapshot(..))
import VM.Bytecode.Reader (readInt32)
import VM.VMStack (stackPop, stackPush)

-- | Implements CALL (Opcode 0x40).
--
-- @details
--   Static call. Reads relative offset, calculates target, and delegates frame setup.
--
instCall :: VirtualMachine ()
instCall :: VirtualMachine ()
instCall = do
    offset <- VirtualMachine Int
readInt32
    vm <- get
    doSnapshot (bytecodeIndex vm + offset) V.empty

-- | Implements TAIL_CALL (Opcode 0x41).
--
-- @details
--   Optimized call. Replaces current frame (no snapshot pushed).
--
instTailCall :: VirtualMachine ()
instTailCall :: VirtualMachine ()
instTailCall = do
    offset <- VirtualMachine Int
readInt32
    vm <- get
    put $ vm { baseVStackIndex = V.length (vStack vm),
        bytecodeIndex = bytecodeIndex vm + offset, env = V.empty }

-- | Helper to dispatch CALL_INDIRECT based on the callable type.
--
-- @args
--   - callee: The value popped from stack (must be VClosure or VFuncPtr).
--
-- @details
--   Extracts the address and environment from the callable and calls 'doSnapshot'.
--
-- @throws
--   Error if 'callee' is not a callable type.
--
dispatchIndirectCall :: VMValue -> VirtualMachine ()
dispatchIndirectCall :: VMValue -> VirtualMachine ()
dispatchIndirectCall (VClosure Int
addr Vector VMValue
caps) = Int -> Vector VMValue -> VirtualMachine ()
doSnapshot Int
addr Vector VMValue
caps
dispatchIndirectCall (VFuncPtr Int
addr) = Int -> Vector VMValue -> VirtualMachine ()
doSnapshot Int
addr Vector VMValue
forall a. Vector a
V.empty
dispatchIndirectCall VMValue
x = [Char] -> VirtualMachine ()
forall a. HasCallStack => [Char] -> a
error ([Char] -> VirtualMachine ()) -> [Char] -> VirtualMachine ()
forall a b. (a -> b) -> a -> b
$ [Char]
"VM Error: Not callable: " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ VMValue -> [Char]
forall a. Show a => a -> [Char]
show VMValue
x

-- | Implements CALL_INDIRECT (Opcode 0x42).
--
-- @details
--   Pops the top value and delegates to 'dispatchIndirectCall'.
--
instCallIndirect :: VirtualMachine ()
instCallIndirect :: VirtualMachine ()
instCallIndirect = do
    callee <- VirtualMachine VMValue
stackPop
    dispatchIndirectCall callee

-- | Helper to finalize the creation of a closure.
--
-- @args
--   - vm: The current VM state.
--   - addr: The function address.
--   - count: The number of variables to capture.
--
-- @details
--   Extracts the last 'count' elements from the stack, truncates the stack,
--   and pushes the new VClosure.
--
pushClosure :: VMState -> Int -> Int -> VirtualMachine ()
pushClosure :: VMState -> Int -> Int -> VirtualMachine ()
pushClosure VMState
vm Int
addr Int
count =
    VMState -> VirtualMachine ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VMState
vm { vStack = V.take start (vStack vm) }) VirtualMachine () -> VirtualMachine () -> VirtualMachine ()
forall a b.
StateT VMState IO a -> StateT VMState IO b -> StateT VMState IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>>
    VMValue -> VirtualMachine ()
stackPush (Int -> Vector VMValue -> VMValue
VClosure Int
addr (Int -> Int -> Vector VMValue -> Vector VMValue
forall a. Int -> Int -> Vector a -> Vector a
V.slice Int
start Int
count (VMState -> Vector VMValue
vStack VMState
vm)))
    where start :: Int
start = (Vector VMValue -> Int
forall a. Vector a -> Int
V.length (VMState -> Vector VMValue
vStack VMState
vm)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
count

-- | Implements MAKE_CLOSURE (Opcode 0x60).
--
-- @details
--   Reads arguments then checks if the stack has enough elements.
--   If yes, calls 'pushClosure'.
--
instMakeClosure :: VirtualMachine ()
instMakeClosure :: VirtualMachine ()
instMakeClosure = do
    offset <- VirtualMachine Int
readInt32
    count <- readInt32
    vm <- get
    let absAddr = VMState -> Int
bytecodeIndex VMState
vm Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
offset
    case V.length (vStack vm) >= count of
        Bool
False -> [Char] -> VirtualMachine ()
forall a. HasCallStack => [Char] -> a
error [Char]
"VM Error: MAKE_CLOSURE Stack Underflow"
        Bool
True -> VMState -> Int -> Int -> VirtualMachine ()
pushClosure VMState
vm Int
absAddr Int
count

-- | Helper to perform the state restoration for RET.
--
-- @args
--   - vm: Current state.
--   - snap: The snapshot to restore.
--   - rest: The remaining snapshot stack.
--   - v: The return value to push.
--
-- @details
--   Restores IP, FP, Env, and cleans the stack before pushing the return value.
--
execReturn :: VMState -> CallSnapshot -> [CallSnapshot] -> VMValue -> Int ->
    VirtualMachine ()
execReturn :: VMState
-> CallSnapshot
-> [CallSnapshot]
-> VMValue
-> Int
-> VirtualMachine ()
execReturn VMState
vm CallSnapshot
snap [CallSnapshot]
rest VMValue
v Int
nArgs = VMState -> VirtualMachine ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (VMState
vm { bytecodeIndex = callbackIndex snap,
    baseVStackIndex = vStackIndex snap, env = vEnv snap, snapshotStack = rest,
    vStack = V.take (baseVStackIndex vm - nArgs) (vStack vm) }) VirtualMachine () -> VirtualMachine () -> VirtualMachine ()
forall a b.
StateT VMState IO a -> StateT VMState IO b -> StateT VMState IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> VMValue -> VirtualMachine ()
stackPush VMValue
v

-- | Implements RET (Opcode 0x43).
--
-- @details
--   1. Pops the return value from the top of the stack.
--   2. Retrieves and removes the top 'CallSnapshot'.
--   3. "Cleans" the stack by truncating it to the start of the current frame.
--      This effectively discards all local variables.
--   4. Restores the Caller's context (IP and FP).
--   5. Pushes the return value back onto the restored stack.
--
-- @throws
--   Error "VM Error: Return called with empty call stack" if executed in global scope.
--
instRet :: VirtualMachine ()
instRet :: VirtualMachine ()
instRet = do
    nArgs <- VirtualMachine Int
readInt32
    retVal <- stackPop
    vm <- get
    case snapshotStack vm of
        [] -> [Char] -> VirtualMachine ()
forall a. HasCallStack => [Char] -> a
error [Char]
"VM Error: Return called with empty call stack"
        (CallSnapshot
snap:[CallSnapshot]
rest) -> VMState
-> CallSnapshot
-> [CallSnapshot]
-> VMValue
-> Int
-> VirtualMachine ()
execReturn VMState
vm CallSnapshot
snap [CallSnapshot]
rest VMValue
retVal Int
nArgs

-- | Implements GET_FUNC_ADDR (Opcode 0x61).
--
-- @details
--   Reads an absolute address (Int) and pushes it as a VFuncPtr.
--   Used to pass functions as values before potentially wrapping them in closures.
--
instGetFuncAddr :: VirtualMachine ()
instGetFuncAddr :: VirtualMachine ()
instGetFuncAddr = do
    addr <- VirtualMachine Int
readInt32
    stackPush (VFuncPtr addr)