-- | Object file generation.
module Language.Elna.Object.ElfCoder
    ( ElfEnvironment(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , UnresolvedRelocation(..)
    , UnresolvedRelocations(..)
    , addHeaderToResult
    , addSectionHeader
    , elfHeaderSize
    , elfObject
    , elfSectionsSize
    , putSectionHeader
    , partitionSymbols
    , 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.Bits (Bits(..))
import Data.ByteString (StrictByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Builder as ByteString.Builder
import Data.Word (Word8)
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 UnresolvedRelocation = UnresolvedRelocation StrictByteString Elf32_Addr Word8
data UnresolvedRelocations =
    UnresolvedRelocations (Vector UnresolvedRelocation) (ElfHeaderResult Elf32_Sym) Elf32_Word

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

partitionSymbols :: ElfHeaderResult Elf32_Sym -> (Vector Elf32_Sym, Vector Elf32_Sym)
partitionSymbols =  Vector.partition go . getField @"sectionHeaders"
  where
    go Elf32_Sym{ st_info } = (st_info .&. 0xf0) == 0

-- | 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

addHeaderToResult :: StrictByteString -> a -> ElfHeaderResult a -> ElfHeaderResult a
addHeaderToResult name newHeader accumulator@ElfHeaderResult{..} = accumulator
    { sectionHeaders = Vector.snoc sectionHeaders newHeader
    , sectionNames = StringTable.append name sectionNames
    }

addSectionHeader :: StrictByteString -> Elf32_Shdr -> ElfWriter ()
addSectionHeader name newHeader = ElfWriter $ modify' modifier
  where
    modifier elfEnvironment@ElfEnvironment{ objectHeaders } = elfEnvironment
        { objectHeaders = addHeaderToResult name newHeader objectHeaders
        }

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