{-
-- EPITECH PROJECT, 2025
-- Glados
-- File description:
-- VM State
-}

module VM.VMState
    ( VMState(..)
    , VirtualMachine
    , createVMState
    , createSnapshot
    , doSnapshot
    ) where

import qualified Data.Map.Strict as Map
import qualified Data.ByteString as BS
import qualified Data.Vector as V
import Control.Monad.State.Strict (StateT, get, put)
import System.IO (Handle, stdin, stdout, stderr)

import VM.VMValue (VMValue)
import VM.CallSnapshot (CallSnapshot(..))

-- | The core state of the Virtual Machine.
--
-- @details
--   This structure holds the entire context of the executing program at any
--   given microsecond. It includes memory (stack, heap/globals), the processor
--   registers (IP, FP), and the code itself.
--
data VMState = VMState
    { VMState -> ByteString
bytecode :: BS.ByteString
      -- ^ The read-only byte array containing the compiled machine code.

    , VMState -> Int
bytecodeIndex :: Int
      -- ^ The current index in 'bytecode'.
      --   When a function is called, this value (plus offset) becomes the 
      --   'callbackIndex' in the pushed CallSnapshot.

    , VMState -> Vector VMValue
vStack :: V.Vector VMValue
      -- ^ The main data stack.
      --   Contains inputs, outputs, and local variables.

    , VMState -> Int
baseVStackIndex :: Int
      -- ^ The index in 'vStack' where the current function's local variables begin.
      --   When a function is called, this value is saved as 'vStackIndex' 
      --   in the CallSnapshot.

    , VMState -> [CallSnapshot]
snapshotStack :: [CallSnapshot]
      -- ^ The stack of saved contexts (formerly callStack).
      --   Used to restore the state upon returning from a function.

    , VMState -> Vector VMValue
env :: V.Vector VMValue
      -- ^ The Current Environment (Closures).
      --   Contains variables captured by the currently executing function.
      --   If the current function is not a closure, this is typically empty.

    , VMState -> Vector VMValue
globalEnv :: V.Vector VMValue
      -- ^ The persistent global environment.
      --   Stores global variables accessible throughout the program life.

    , VMState -> Bool
isRunning :: Bool
      -- ^ Execution flag. If False, the VM loop terminates.

    , VMState -> Bool
isDebug :: Bool
      -- ^ Enable tracing

    , VMState -> Map Int Handle
vmFds :: Map.Map Int Handle
      -- ^ Table of file descriptors (Kernel)

    , VMState -> Int
nextFd :: Int
      -- ^ Next free fd
    }

-- | The Monad Transformer stack used for the VM.
--
-- @details
--   We use 'StateT' to carry the 'VMState' and 'IO' to allow for
--   side effects like printing to the console (PRINT instruction).
type VirtualMachine a = StateT VMState IO a

-- | Creates the initial state of the VM for a given program.
--
-- @args
--   - code: The compiled binary code to execute.
--
-- @return
--   A pristine 'VMState' ready for execution.
--
createVMState :: BS.ByteString -> Bool -> VMState
createVMState :: ByteString -> Bool -> VMState
createVMState ByteString
code Bool
debug = VMState
    { bytecode :: ByteString
bytecode = ByteString
code, bytecodeIndex :: Int
bytecodeIndex = Int
0
    , vStack :: Vector VMValue
vStack = Vector VMValue
forall a. Vector a
V.empty, baseVStackIndex :: Int
baseVStackIndex = Int
0
    , snapshotStack :: [CallSnapshot]
snapshotStack = [], env :: Vector VMValue
env = Vector VMValue
forall a. Vector a
V.empty
    , globalEnv :: Vector VMValue
globalEnv = Int -> VMValue -> Vector VMValue
forall a. Int -> a -> Vector a
V.replicate Int
1024 (VMValue
forall a. HasCallStack => a
undefined)
    , isRunning :: Bool
isRunning = Bool
True, isDebug :: Bool
isDebug = Bool
debug
    , vmFds :: Map Int Handle
vmFds = [(Int, Handle)] -> Map Int Handle
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Int
0, Handle
stdin), (Int
1, Handle
stdout), (Int
2, Handle
stderr)]
    , nextFd :: Int
nextFd = Int
3 }

-- | Creates a snapshot of the current VM context (IP, FP, Env).
--
-- @args
--   - vm: The current VM state.
--
-- @details
--   Captures the current Instruction Pointer, Frame Pointer, and Closure Environment.
--   Used by CALL instructions to save the state before jumping.
--
-- @return
--   A 'CallSnapshot' containing the saved context.
--
createSnapshot :: VMState -> CallSnapshot
createSnapshot :: VMState -> CallSnapshot
createSnapshot VMState
vm = CallSnapshot
    { callbackIndex :: Int
callbackIndex = VMState -> Int
bytecodeIndex VMState
vm
    , vStackIndex :: Int
vStackIndex = VMState -> Int
baseVStackIndex VMState
vm
    , vEnv :: Vector VMValue
vEnv = VMState -> Vector VMValue
env VMState
vm }

-- | Sets up a new stack frame for a function call (Shared Logic).
--
-- @args
--   - idx: The absolute address (index) to jump to.
--   - newEnv: The closure environment to load (empty for static calls).
--
-- @details
--   1. Creates and pushes a snapshot of the current state.
--   2. Updates FP (baseVStackIndex) and IP (bytecodeIndex).
--   3. Sets the new environment.
--
-- @return
--   Unit. State is modified.
--
doSnapshot :: Int -> V.Vector VMValue -> VirtualMachine ()
doSnapshot :: Int -> Vector VMValue -> VirtualMachine ()
doSnapshot Int
idx Vector VMValue
nEnv = do
    vm <- StateT VMState IO VMState
forall s (m :: * -> *). MonadState s m => m s
get
    put $ vm { snapshotStack = createSnapshot vm : snapshotStack vm
             , baseVStackIndex = V.length (vStack vm)
             , bytecodeIndex = idx
             , env = nEnv }