-- | Object file generation.
module Language.Elna.Object.ElfCoder
    ( ElfEnvironment(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , elfHeaderSize
    , addSectionHeader
    , elfObject
    , elfSectionsSize
    , putSectionHeader
    , module Language.Elna.Object.Elf
    ) where

import Control.Exception (throwIO)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.State (StateT, runStateT, modify', gets)
import Data.ByteString (StrictByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as ByteString.Builder
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
import Data.Foldable (traverse_)
import Language.Elna.Object.Elf
import Language.Elna.Object.StringTable (StringTable)
import qualified Language.Elna.Object.StringTable as StringTable
import GHC.Records (HasField(..))

data ElfEnvironment = ElfEnvironment
    { objectHeaders :: ElfHeaderResult Elf32_Shdr
    , objectHandle :: Handle
    }

newtype ElfWriter a = ElfWriter
    { runElfWriter :: StateT ElfEnvironment IO a
    }

data ElfHeaderResult a = ElfHeaderResult
    { sectionNames :: StringTable
    , sectionHeaders :: Vector a
    } deriving Eq

instance Functor ElfWriter
  where
    fmap f (ElfWriter x) = ElfWriter $ f <$> x

instance Applicative ElfWriter
  where
    pure = ElfWriter . pure
    (ElfWriter f) <*> (ElfWriter x) = ElfWriter $ f <*> x

instance Monad ElfWriter
  where
    (ElfWriter x) >>= f = ElfWriter $ x >>= (runElfWriter . f)

instance MonadIO ElfWriter
  where
    liftIO = ElfWriter . liftIO

-- | ELF header size.
elfHeaderSize :: Elf32_Off
elfHeaderSize = 52

-- | Calculates the size of all sections based on the 'sh_size' in the given
-- headers and adds 'elfHeaderSize' to it.
elfSectionsSize :: Vector Elf32_Shdr -> Elf32_Off
elfSectionsSize = (elfHeaderSize +)
    . Vector.foldr ((+) . sh_size) 0

addSectionHeader :: StrictByteString -> Elf32_Shdr -> ElfWriter ()
addSectionHeader name newHeader = ElfWriter $ modify' modifier
  where
    modifier elfEnvironment@ElfEnvironment{ objectHeaders } =
        let ElfHeaderResult{..} = objectHeaders
         in elfEnvironment
            { objectHeaders = ElfHeaderResult
                { sectionHeaders = Vector.snoc sectionHeaders newHeader
                , sectionNames = StringTable.append name sectionNames
                }
            }

putSectionHeader :: StrictByteString -> Elf32_Shdr -> StrictByteString -> ElfWriter ()
putSectionHeader name newHeader encoded = do
    objectHandle' <- ElfWriter $ gets $ getField @"objectHandle"
    liftIO $ ByteString.hPut objectHandle' encoded
    addSectionHeader name newHeader

-- Writes an ELF object to the provided file path. The callback writes the
-- sections, collects headers for those sections and returns the ELF header.
elfObject :: FilePath -> ElfWriter Elf32_Ehdr -> IO ()
elfObject outFile putContents = withFile outFile WriteMode withObjectFile
  where
    withObjectFile objectHandle
        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
        >> putContents' objectHandle
        >>= uncurry afterContents
    putContents' objectHandle
        = flip runStateT (initialState objectHandle)
        $ runElfWriter putContents
    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
        }
    initialState objectHandle = ElfEnvironment
        { objectHeaders = ElfHeaderResult
            { sectionHeaders = Vector.singleton zeroHeader
            , sectionNames = mempty
            }
        , objectHandle = objectHandle
        }
    afterContents header ElfEnvironment{ objectHeaders = ElfHeaderResult{..}, ..} =
        let hPutBuilder = ByteString.Builder.hPutBuilder objectHandle
            writeSectionHeaders byteOrder' =
                traverse_ (hPutBuilder . elf32Shdr byteOrder') sectionHeaders
         in either throwIO pure (byteOrder (e_ident header))
            >>= writeSectionHeaders
            >> either throwIO (putHeaders objectHandle) (elf32Ehdr header)
    putHeaders objectHandle encodedHeader
        = hSeek objectHandle AbsoluteSeek 0
        >> ByteString.Builder.hPutBuilder objectHandle encodedHeader