{-
-- EPITECH PROJECT, 2026
-- Glados
-- File description:
-- Assembler.hs
-}

module Compiler.ASM.Assembler
    ( assemble
    ) where

import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString as BS
import Data.Text (Text, pack)
import Data.Int (Int32)
import qualified Data.Set as Set
import Data.Foldable (traverse_)

import Compiler.PsInstruction (PsInstruction)
import Compiler.Instruction (Instruction(..), instructionSize)
import Compiler.Bytecode.Serializer (serializeInstruction)
import Compiler.ResolveLabels.ResolveLabels (resolveLabels)
import Compiler.Bytecode.Encoder (encodeInt32BE)

magicBytes :: B.Builder
magicBytes :: Builder
magicBytes = Word8 -> Builder
B.word8 Word8
0x47 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0x4C Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0x41 Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Word8 -> Builder
B.word8 Word8
0x44 -- GLAD

versionBytes :: B.Builder
versionBytes :: Builder
versionBytes = Word8 -> Builder
B.word8 Word8
0x02

flagBytes :: B.Builder
flagBytes :: Builder
flagBytes = Word8 -> Builder
B.word8 Word8
0x00

-- | Main assembly function.
--
-- @details
--   1. Resolves labels (PsInstruction -> Instruction) using ResolveLabels module.
--   2. Validates jump offsets (bounds checking).
--   3. Serializes instructions to ByteString with Header.
--
assemble :: [PsInstruction] -> Either Text BS.ByteString
assemble :: [PsInstruction] -> Either Text ByteString
assemble [PsInstruction]
ps = do
    insts <- [PsInstruction] -> Either Text [Instruction]
resolveLabels [PsInstruction]
ps

    let totalSize = [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ((Instruction -> Int) -> [Instruction] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Instruction -> Int
instructionSize [Instruction]
insts)
    validateJumps insts totalSize

    return $ finalizeBytecode insts

-- | Serialize the resolved instructions and prepend the header.
--
finalizeBytecode :: [Instruction] -> BS.ByteString
finalizeBytecode :: [Instruction] -> ByteString
finalizeBytecode [Instruction]
insts =
    let body :: LazyByteString
body = Builder -> LazyByteString
B.toLazyByteString (Builder -> LazyByteString) -> Builder -> LazyByteString
forall a b. (a -> b) -> a -> b
$ (Instruction -> Builder) -> [Instruction] -> Builder
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Instruction -> Builder
serializeInstruction [Instruction]
insts
        header :: Builder
header = Int32 -> Builder
buildHeader (Int64 -> Int32
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64 -> Int32) -> Int64 -> Int32
forall a b. (a -> b) -> a -> b
$ LazyByteString -> Int64
BL.length LazyByteString
body) in
    LazyByteString -> ByteString
BL.toStrict (LazyByteString -> ByteString)
-> (Builder -> LazyByteString) -> Builder -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyByteString
B.toLazyByteString (Builder -> ByteString) -> Builder -> ByteString
forall a b. (a -> b) -> a -> b
$ Builder
header Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> LazyByteString -> Builder
B.lazyByteString LazyByteString
body

buildHeader :: Int32 -> B.Builder
buildHeader :: Int32 -> Builder
buildHeader Int32
size =
    Builder
magicBytes Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
versionBytes Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
flagBytes Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Int32 -> Builder
encodeInt32BE Int32
size

-- | Ensures all jumps point to valid instruction boundaries within the code.
--
validateJumps :: [Instruction] -> Int -> Either Text ()
validateJumps :: [Instruction] -> Int -> Either Text ()
validateJumps [Instruction]
insts Int
size =
    let offs :: [Int]
offs = (Int -> Instruction -> Int) -> Int -> [Instruction] -> [Int]
forall b a. (b -> a -> b) -> b -> [a] -> [b]
scanl (\Int
off Instruction
i -> Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Instruction -> Int
instructionSize Instruction
i) Int
0 [Instruction]
insts
        validOffsets :: Set Int
validOffsets = [Int] -> Set Int
forall a. Ord a => [a] -> Set a
Set.fromList [Int]
offs in
    ((Int, Instruction) -> Either Text ())
-> [(Int, Instruction)] -> Either Text ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ (Int -> Set Int -> (Int, Instruction) -> Either Text ()
checkJump Int
size Set Int
validOffsets) ([Int] -> [Instruction] -> [(Int, Instruction)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Int]
offs [Instruction]
insts)

-- | Checks a single instruction if it's a Jump.
--
checkJump :: Int -> Set.Set Int -> (Int, Instruction) -> Either Text ()
checkJump :: Int -> Set Int -> (Int, Instruction) -> Either Text ()
checkJump Int
size Set Int
validSet (Int
currentOffset, Instruction
inst) =
    case Int -> Int -> Instruction -> Maybe Int
getJumpTarget Int
currentOffset (Instruction -> Int
instructionSize Instruction
inst) Instruction
inst of
        Maybe Int
Nothing -> () -> Either Text ()
forall a b. b -> Either a b
Right ()
        Just Int
target -> Int -> Set Int -> Int -> Either Text ()
verifyTarget Int
size Set Int
validSet Int
target

verifyTarget :: Int -> Set.Set Int -> Int -> Either Text ()
verifyTarget :: Int -> Set Int -> Int -> Either Text ()
verifyTarget Int
size Set Int
validSet Int
t = case Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0 Bool -> Bool -> Bool
&& Int
t Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
size of
    Bool
True -> case Int -> Set Int -> Bool
forall a. Ord a => a -> Set a -> Bool
Set.member Int
t Set Int
validSet of
        Bool
True -> () -> Either Text ()
forall a b. b -> Either a b
Right ()
        Bool
False -> Text -> Either Text ()
forall a b. a -> Either a b
Left (Text -> Either Text ()) -> Text -> Either Text ()
forall a b. (a -> b) -> a -> b
$ String -> Text
pack String
"Jump target not aligned (mid-instruction): " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>
            String -> Text
pack (Int -> String
forall a. Show a => a -> String
show Int
t)
    Bool
False -> Text -> Either Text ()
forall a b. a -> Either a b
Left (Text -> Either Text ()) -> Text -> Either Text ()
forall a b. (a -> b) -> a -> b
$ String -> Text
pack String
"Jump out of bounds: " Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> String -> Text
pack (Int -> String
forall a. Show a => a -> String
show Int
t)

-- | Helper to calculate absolute target of relative jumps.
--
getJumpTarget :: Int -> Int -> Instruction -> Maybe Int
getJumpTarget :: Int -> Int -> Instruction -> Maybe Int
getJumpTarget Int
off Int
sz (Jump Int
rel)          = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
off Int
sz (JumpIfFalse Int
rel)   = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
off Int
sz (JumpIfTrue Int
rel)    = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
off Int
sz (Call Int
rel)          = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
off Int
sz (TailCall Int
rel)      = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
off Int
sz (MakeClosure Int
rel Int
_) = Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
off Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
sz Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
rel)
getJumpTarget Int
_ Int
_ Instruction
_ = Maybe Int
forall a. Maybe a
Nothing