-- | Writer assembler to an object file.
module Language.Elna.PrinterWriter
    ( riscv32Elf
    , riscv32Header
    ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as ByteString.Builder
import qualified Data.ByteString.Lazy as LazyByteString
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import Language.Elna.Object.Elf
    ( ByteOrder(..)
    , Elf32_Ehdr(..)
    , Elf32_Half
    , Elf32_Sym(..)
    , ElfMachine(..)
    , ElfType(..)
    , ElfVersion(..)
    , ElfIdentification(..)
    , ElfClass(..)
    , ElfData(..)
    , Elf32_Shdr(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , elf32Sym
    , elfHeaderSize
    , elfSectionsSize
    , stInfo
    )
import System.IO (Handle)
import qualified Language.Elna.Architecture.RiscV as RiscV

data ElfHeaderResult a = ElfHeaderResult ByteString (Vector a)

riscv32Elf :: Vector RiscV.Instruction -> Handle -> IO (Vector Elf32_Shdr)
riscv32Elf code objectHandle =
    let zeroHeader = Elf32_Shdr
            { sh_type = SHT_NULL
            , sh_size = 0
            , sh_offset = 0
            , sh_name = 0
            , sh_link = 0
            , sh_info = 0
            , sh_flags = 0
            , sh_entsize = 0
            , sh_addralign = 0
            , sh_addr = 0
            }
     in text (ElfHeaderResult "\0" $ Vector.singleton zeroHeader)
        >>= shstrtab
        >>= finalize
  where
    finalize (ElfHeaderResult _ headers) = pure headers
    shstrtab (ElfHeaderResult names headers) = do
        let stringTable = names <> ".shstrtab\0"
            nextHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle stringTable
        pure $ ElfHeaderResult stringTable
            $ Vector.snoc headers nextHeader
    strtab stringTable (ElfHeaderResult names headers) = do
        let newHeader = Elf32_Shdr
                { sh_type = SHT_STRTAB
                , sh_size = fromIntegral $ ByteString.length stringTable
                , sh_offset = elfSectionsSize headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle stringTable
        pure $ ElfHeaderResult (names <> ".strtab\0")
            $ Vector.snoc headers newHeader
    symtab strtabIndex entries (ElfHeaderResult names headers) = do
        let encoded = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ foldMap (elf32Sym LSB) entries 
            newHeader = Elf32_Shdr
                { sh_type = SHT_SYMTAB
                , sh_size = fromIntegral $ ByteString.length encoded
                , sh_offset = elfSectionsSize headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = strtabIndex
                , sh_info = 1
                , sh_flags = 0
                , sh_entsize = 16
                , sh_addralign = 0
                , sh_addr = 0
                }
        ByteString.hPut objectHandle encoded
        pure $ ElfHeaderResult (names <> ".symtab\0")
            $ Vector.snoc headers newHeader
    text (ElfHeaderResult names headers) = do
        let textTabIndex = fromIntegral $ Vector.length headers
            strtabIndex = fromIntegral $ textTabIndex + 2
        ElfHeaderResult stringTable  entries <- symbolEntry textTabIndex code
            $ ElfHeaderResult "\0"
            $ Vector.singleton
            $ Elf32_Sym
                { st_value = 0
                , st_size = 0
                , st_shndx = 0
                , st_other = 0
                , st_name = 0
                , st_info = 0
                }
        let newHeader = Elf32_Shdr
                { sh_type = SHT_PROGBITS
                , sh_size = fromIntegral $ foldr ((+) . st_size) 0 entries
                , sh_offset = elfSectionsSize headers
                , sh_name = fromIntegral $ ByteString.length names
                , sh_link = 0
                , sh_info = 0
                , sh_flags = 0b110
                , sh_entsize = 0
                , sh_addralign = 0
                , sh_addr = 0
                }
            newResult = ElfHeaderResult (names <> ".text\0")
                $ Vector.snoc headers newHeader
        symtab strtabIndex entries newResult
            >>= strtab stringTable
    symbolEntry :: Elf32_Half -> Vector RiscV.Instruction -> ElfHeaderResult Elf32_Sym -> IO (ElfHeaderResult Elf32_Sym)
    symbolEntry shndx instructions (ElfHeaderResult names entries) = do
        let encoded = LazyByteString.toStrict
                $ ByteString.Builder.toLazyByteString
                $ foldMap RiscV.instruction instructions
            newEntry = Elf32_Sym
                { st_value = 0
                , st_size = fromIntegral $ ByteString.length encoded
                , st_shndx = shndx
                , st_other = 0
                , st_name = fromIntegral $ ByteString.length names
                , st_info = stInfo STB_GLOBAL STT_FUNC
                }
        ByteString.hPut objectHandle encoded
        pure $ ElfHeaderResult (names <> "_start\0")
            $ Vector.snoc entries newEntry

riscv32Header :: Elf32_Ehdr
riscv32Header = Elf32_Ehdr
    { e_version = EV_CURRENT
    , e_type = ET_REL
    , e_shstrndx = 2 -- String table. SHN_UNDEF
    , e_shoff = 0
    , e_shnum = 0
    , 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
    }