-- | Writer assembler to an object file.
module Language.Elna.RiscV.ElfWriter
    ( riscv32Elf
    ) where

import Data.ByteString (StrictByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as ByteString.Builder
import Data.ByteString.Lazy (LazyByteString)
import qualified Data.ByteString.Lazy as LazyByteString
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Language.Elna.Object.ElfCoder
    ( ByteOrder(..)
    , Elf32_Ehdr(..)
    , Elf32_Half
    , Elf32_Word
    , Elf32_Sym(..)
    , ElfMachine(..)
    , ElfType(..)
    , ElfVersion(..)
    , ElfIdentification(..)
    , ElfClass(..)
    , ElfData(..)
    , Elf32_Shdr(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , Elf32_Rel(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , ElfEnvironment(..)
    , UnresolvedRelocation(..)
    , UnresolvedRelocations(..)
    , addHeaderToResult
    , addSectionHeader
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    , rInfo
    , elf32Rel
    , shfInfoLink
    , partitionSymbols
    , putSectionHeader
    )
import qualified Language.Elna.Architecture.RiscV as RiscV
import qualified Data.Text.Encoding as Text
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.State (get, gets)
import Language.Elna.RiscV.CodeGenerator (Directive(..), Statement(..))
import Language.Elna.Object.StringTable (StringTable)
import qualified Language.Elna.Object.StringTable as StringTable
import Data.HashSet (HashSet)
import qualified Data.HashSet as HashSet
import GHC.Records (HasField(..))

data TextAccumulator = TextAccumulator
    { encodedAccumulator :: LazyByteString
    , relocationAccumulator :: Vector UnresolvedRelocation
    , symbolAccumulator :: ElfHeaderResult Elf32_Sym
    , definitionAccumulator :: HashSet StrictByteString
    }

riscv32Elf :: Vector Statement -> ElfWriter Elf32_Ehdr
riscv32Elf code = text code
    >>= symtab
    >>= uncurry symrel
    >>= strtab
    >> shstrtab
    >>= riscv32Header
  where
    riscv32Header :: Elf32_Half -> ElfWriter Elf32_Ehdr
    riscv32Header shstrndx = do
        ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
        pure $ Elf32_Ehdr
            { e_version = EV_CURRENT
            , e_type = ET_REL
            , e_shstrndx = shstrndx
            , e_shoff = elfSectionsSize sectionHeaders
            , e_shnum = fromIntegral (Vector.length sectionHeaders)
            , e_shentsize = 40
            , e_phoff = 0
            , e_phnum = 0
            , e_phentsize = 32
            , e_machine = EM_RISCV
            , e_ident = ElfIdentification ELFCLASS32 ELFDATA2LSB
            , e_flags = 0x4 -- EF_RISCV_FLOAT_ABI_DOUBLE
            , e_entry = 0
            , e_ehsize = fromIntegral elfHeaderSize
            }

text :: Vector Statement -> ElfWriter UnresolvedRelocations
text code = do
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
    let textTabIndex = fromIntegral $ Vector.length sectionHeaders
        initialHeaders = ElfHeaderResult mempty
            $ Vector.singleton
            $ Elf32_Sym
                { st_value = 0
                , st_size = 0
                , st_shndx = 0
                , st_other = 0
                , st_name = 0
                , st_info = 0
                }
        TextAccumulator{..} = encodeFunctions textTabIndex code
            $ TextAccumulator
                { encodedAccumulator = mempty 
                , relocationAccumulator = Vector.empty
                , symbolAccumulator = initialHeaders
                , definitionAccumulator = HashSet.empty
                }
        size = fromIntegral $ LazyByteString.length encodedAccumulator
        newHeader = Elf32_Shdr
            { sh_type = SHT_PROGBITS
            , sh_size = size
            , sh_offset = elfSectionsSize sectionHeaders
            , sh_name = StringTable.size sectionNames
            , sh_link = 0
            , sh_info = 0
            , sh_flags = 0b110
            , sh_entsize = 0
            , sh_addralign = 4
            , sh_addr = 0
            }
    putSectionHeader ".text" newHeader $ LazyByteString.toStrict encodedAccumulator
    let filterPredicate :: StrictByteString -> Bool
        filterPredicate = not
            . (`StringTable.elem` getField @"sectionNames" symbolAccumulator)
        symbolResult = HashSet.foldl' encodeEmptyDefinitions symbolAccumulator
            $ HashSet.filter filterPredicate definitionAccumulator
    pure $ UnresolvedRelocations relocationAccumulator symbolResult
        $ fromIntegral $ Vector.length sectionHeaders
  where
    encodeEmptyDefinitions (ElfHeaderResult names entries) definition =
        let nextEntry = Elf32_Sym
                { st_value = 0
                , st_size = 0
                , st_shndx = 0
                , st_other = 0
                , st_name = StringTable.size names
                , st_info = stInfo STB_GLOBAL STT_FUNC
                }
         in ElfHeaderResult (StringTable.append definition names)
            $ Vector.snoc entries nextEntry
    encodeFunctions shndx instructions textAccumulator
        | Just (instruction, rest) <- Vector.uncons instructions =
            case instruction of
                Instruction _ ->
                    let (textAccumulator', rest') = encodeInstructions shndx (textAccumulator, instructions)
                     in encodeFunctions shndx rest' textAccumulator'
                JumpLabel labelName directives ->
                    let (TextAccumulator{..}, rest') =
                            encodeInstructions shndx (textAccumulator, rest)
                        newEntry = Elf32_Sym
                            { st_value = fromIntegral
                                $ LazyByteString.length
                                $ getField @"encodedAccumulator" textAccumulator
                            , st_size = fromIntegral $ LazyByteString.length encodedAccumulator
                            , st_shndx = shndx
                            , st_other = 0
                            , st_name = StringTable.size $ getField @"sectionNames" symbolAccumulator
                            , st_info = stInfo (directivesBinding directives) STT_FUNC
                            }
                     in encodeFunctions shndx rest'
                        $ TextAccumulator
                            { encodedAccumulator = encodedAccumulator
                            , relocationAccumulator = relocationAccumulator
                            , symbolAccumulator =
                                addHeaderToResult (Text.encodeUtf8 labelName) newEntry symbolAccumulator
                            , definitionAccumulator = definitionAccumulator
                            }
        | otherwise = textAccumulator
    directivesBinding directives
        | GlobalDirective `elem` directives = STB_GLOBAL
        | otherwise = STB_LOCAL
    encodeInstructions shndx (TextAccumulator encoded relocations symbolResult definitions, instructions)
        | Just (Instruction instruction, rest) <- Vector.uncons instructions =
            let offset = fromIntegral $ LazyByteString.length encoded
                unresolvedRelocation = case instruction of
                    RiscV.RelocatableInstruction _ instructionType
                        | RiscV.RHigher20 _ symbolName <- instructionType
                            -> Just -- R_RISCV_HI20
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 26
                        | RiscV.RLower12I _ _ _ symbolName <- instructionType
                            -> Just -- R_RISCV_LO12_I
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 27
                        | RiscV.RLower12S symbolName _ _ _ <- instructionType
                            -> Just -- R_RISCV_LO12_S
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 28
                        | RiscV.RBranch symbolName _ _ _ <- instructionType
                            -> Just -- R_RISCV_BRANCH
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 16
                        | RiscV.RJal _ symbolName <- instructionType
                            -> Just -- R_RISCV_JAL
                            $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 17
                    RiscV.CallInstruction symbolName
                        -> Just -- R_RISCV_CALL_PLT
                        $ UnresolvedRelocation (Text.encodeUtf8 symbolName) offset 19
                    RiscV.BaseInstruction _ _ -> Nothing
                chunk = ByteString.Builder.toLazyByteString
                    $ RiscV.instruction instruction
                result = TextAccumulator
                    (encoded <> chunk)
                    (maybe relocations (Vector.snoc relocations) unresolvedRelocation)
                    symbolResult
                    (addDefinition unresolvedRelocation definitions)
             in encodeInstructions shndx (result, rest)
        | Just (JumpLabel labelName directives , rest) <- Vector.uncons instructions
        , FunctionDirective `notElem` directives =
            let newEntry = Elf32_Sym
                    { st_value = fromIntegral $ LazyByteString.length encoded
                    , st_size = 0
                    , st_shndx = shndx
                    , st_other = 0
                    , st_name = StringTable.size $ getField @"sectionNames" symbolResult
                    , st_info = stInfo (directivesBinding directives) STT_NOTYPE
                    }
                result = TextAccumulator
                    encoded
                    relocations
                    (addHeaderToResult (Text.encodeUtf8 labelName) newEntry symbolResult)
                    definitions
                in encodeInstructions shndx (result, rest)
        | otherwise = (TextAccumulator encoded relocations symbolResult definitions, instructions)
    addDefinition (Just (UnresolvedRelocation symbolName _ _)) =
        HashSet.insert symbolName
    addDefinition Nothing = id

shstrtab :: ElfWriter Elf32_Half
shstrtab = do
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
    let stringTable = ".shstrtab"
        currentNamesSize = StringTable.size sectionNames
        nextHeader = Elf32_Shdr
            { sh_type = SHT_STRTAB
            , sh_size = currentNamesSize -- Adding trailing null character.
                + fromIntegral (succ $ ByteString.length stringTable)
            , sh_offset = elfSectionsSize sectionHeaders
            , sh_name = currentNamesSize
            , sh_link = 0
            , sh_info = 0
            , sh_flags = 0
            , sh_entsize = 0
            , sh_addralign = 1
            , sh_addr = 0
            }
    addSectionHeader stringTable nextHeader

    ElfEnvironment{..} <- ElfWriter get
    liftIO $ ByteString.hPut objectHandle
        $ StringTable.encode
        $ getField @"sectionNames" objectHeaders
    pure $ fromIntegral $ Vector.length sectionHeaders

symtab :: UnresolvedRelocations -> ElfWriter (Elf32_Word, UnresolvedRelocations)
symtab (UnresolvedRelocations relocationList symbolResult index) = do 
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
    let (localSymbols, globalSymbols) = partitionSymbols symbolResult
        sortedSymbols = localSymbols <> globalSymbols
        sortedResult = symbolResult{ sectionHeaders = sortedSymbols }
        encodedSymbols = LazyByteString.toStrict
            $ ByteString.Builder.toLazyByteString
            $ foldMap (elf32Sym LSB) sortedSymbols
        symHeader = Elf32_Shdr
            { sh_type = SHT_SYMTAB
            , sh_size = fromIntegral $ ByteString.length encodedSymbols
            , sh_offset = elfSectionsSize sectionHeaders
            , sh_name = StringTable.size sectionNames
            , sh_link = fromIntegral $ Vector.length sectionHeaders + 2
            , sh_info = fromIntegral $ Vector.length localSymbols
            , sh_flags = 0
            , sh_entsize = 16
            , sh_addralign = 4
            , sh_addr = 0
            }
    putSectionHeader ".symtab" symHeader encodedSymbols
    pure
        ( fromIntegral $ Vector.length sectionHeaders
        , UnresolvedRelocations relocationList sortedResult index
        )

symrel :: Elf32_Word -> UnresolvedRelocations -> ElfWriter StringTable
symrel sectionHeadersLength relocations = do
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"

    let UnresolvedRelocations relocationList symbols index = relocations
        encodedRelocations = LazyByteString.toStrict
            $ ByteString.Builder.toLazyByteString
            $ Vector.foldMap (either (const mempty) (elf32Rel LSB))
            $ resolveRelocation symbols <$> relocationList
        relHeader = Elf32_Shdr
            { sh_type = SHT_REL
            , sh_size = fromIntegral $ ByteString.length encodedRelocations
            , sh_offset = elfSectionsSize sectionHeaders
            , sh_name = StringTable.size sectionNames
            , sh_link = sectionHeadersLength
            , sh_info = index
            , sh_flags = shfInfoLink
            , sh_entsize = 8
            , sh_addralign = 4
            , sh_addr = 0
            }
    putSectionHeader ".rel.text" relHeader encodedRelocations
    pure $ getField @"sectionNames" symbols
  where
    takeStringZ stringTable Elf32_Sym{ st_name }
        = StringTable.index st_name stringTable
    resolveRelocation (ElfHeaderResult stringTable entries) unresolvedRelocation
        | UnresolvedRelocation symbolName offset type' <- unresolvedRelocation
        , Just entry <- Vector.findIndex ((== symbolName) . takeStringZ stringTable) entries =
            Right $ Elf32_Rel
                { r_offset = offset
                , r_info = rInfo (fromIntegral entry) type'
                }
        | otherwise = Left unresolvedRelocation

strtab :: StringTable -> ElfWriter ()
strtab stringTable = do
    ElfHeaderResult{..} <- ElfWriter $ gets $ getField @"objectHeaders"
    let strHeader = Elf32_Shdr
            { sh_type = SHT_STRTAB
            , sh_size = StringTable.size stringTable
            , sh_offset = elfSectionsSize sectionHeaders
            , sh_name = StringTable.size sectionNames
            , sh_link = 0
            , sh_info = 0
            , sh_flags = 0
            , sh_entsize = 0
            , sh_addralign = 1
            , sh_addr = 0
            }
    putSectionHeader ".strtab" strHeader $ StringTable.encode stringTable