module Language.Elna.Object.Elf
    ( ByteOrder(..)
    , Elf32_Addr
    , Elf32_Off
    , Elf32_Half
    , Elf32_Word
    , Elf32_Sword
    , Elf32_Ehdr(..)
    , Elf32_Rel(..)
    , Elf32_Rela(..)
    , Elf32_Shdr(..)
    , Elf32_Sym(..)
    , ElfEncodingError(..)
    , ElfIdentification(..)
    , ElfMachine(..)
    , ElfVersion(..)
    , ElfClass(..)
    , ElfData(..)
    , ElfType(..)
    , ElfSectionType(..)
    , ElfSymbolBinding(..)
    , ElfSymbolType(..)
    , ElfWriter(..)
    , ElfHeaderResult(..)
    , addSectionHeader
    , elf32Addr
    , elf32Half
    , elf32Off
    , elf32Shdr
    , elf32Sword
    , elf32Word
    , elf32Ehdr
    , elf32Rel
    , elf32Rela
    , elf32Sym
    , elfHeaderSize
    , elfIdentification
    , elfObject
    , elfSectionsSize
    , rInfo
    , shfWrite
    , shfAlloc
    , shfExecinstr
    , shfMascproc
    , shfInfoLink
    , stInfo
    ) where

import Control.Exception (Exception(..), throwIO)
import Data.Bits (Bits(..))
import qualified Data.ByteString.Builder as ByteString.Builder
import Data.Int (Int32)
import Data.Word (Word8, Word16, Word32)
import qualified Data.ByteString as ByteString
import Data.Vector (Vector)
import qualified Data.Vector as Vector
import System.IO (Handle, IOMode(..), SeekMode(..), hSeek, withFile)
import Data.Foldable (traverse_)
import Control.Monad.Trans.State (StateT, runStateT, modify')
import Control.Monad.IO.Class (MonadIO(..))
import Data.ByteString (ByteString)

-- * Data types.

type Elf32_Addr = Word32 -- ^ Unsigned program address.
type Elf32_Half = Word16 -- ^ Unsigned medium integer.
type Elf32_Off = Word32 -- ^ Unsigned file offset.
type Elf32_Sword = Int32 -- ^ Signed large integer.
type Elf32_Word = Word32 -- ^ Unsigned large integer.

data ElfClass
    = ELFCLASSNONE -- ^ Invalid class.
    | ELFCLASS32 -- ^ 32-bit objects.
    | ELFCLASS64 -- ^ 64-bit objects.
    deriving Eq

instance Show ElfClass
  where
    show ELFCLASSNONE = "ELFCLASSNONE"
    show ELFCLASS32 = "ELFCLASS32"
    show ELFCLASS64 = "ELFCLASS64"

instance Enum ElfClass
  where
    toEnum 0 = ELFCLASSNONE
    toEnum 1 = ELFCLASS32
    toEnum 2 = ELFCLASS64
    toEnum _ = error "Unknown Elf class"
    fromEnum ELFCLASSNONE = 0
    fromEnum ELFCLASS32 = 1
    fromEnum ELFCLASS64 = 1

-- | Data encoding.
data ElfData
    = ELFDATANONE
    | ELFDATA2LSB
    | ELFDATA2MSB
    deriving Eq

instance Show ElfData
  where
    show ELFDATANONE = "ELFDATANONE"
    show ELFDATA2LSB = "ELFDATA2LSB"
    show ELFDATA2MSB = "ELFDATA2MSB"

instance Enum ElfData
  where
    toEnum 0 = ELFDATANONE
    toEnum 1 = ELFDATA2LSB
    toEnum 2 = ELFDATA2MSB
    toEnum _ = error "Unknown elf data"
    fromEnum ELFDATANONE = 0
    fromEnum ELFDATA2LSB = 1
    fromEnum ELFDATA2MSB = 2

data ElfIdentification = ElfIdentification ElfClass ElfData
    deriving Eq

-- | ELF header.
data Elf32_Ehdr = Elf32_Ehdr
    { e_ident :: ElfIdentification
    , e_type :: ElfType
    , e_machine :: ElfMachine
    , e_version :: ElfVersion
    , e_entry :: Elf32_Addr
    , e_phoff :: Elf32_Off
    , e_shoff :: Elf32_Off
    , e_flags :: Elf32_Word
    , e_ehsize :: Elf32_Half
    , e_phentsize :: Elf32_Half
    , e_phnum :: Elf32_Half
    , e_shentsize :: Elf32_Half
    , e_shnum :: Elf32_Half
    , e_shstrndx :: Elf32_Half
    } deriving Eq

-- | Section header.
data Elf32_Shdr = Elf32_Shdr
    { sh_name :: Elf32_Word
    , sh_type :: ElfSectionType
    , sh_flags :: Elf32_Word
    , sh_addr :: Elf32_Addr
    , sh_offset :: Elf32_Off
    , sh_size :: Elf32_Word
    , sh_link :: Elf32_Word
    , sh_info :: Elf32_Word
    , sh_addralign :: Elf32_Word
    , sh_entsize :: Elf32_Word
    } deriving Eq

data ElfMachine
    = ElfMachine Elf32_Half
    | EM_NONE -- ^ No machine.
    | EM_M32 -- ^ AT&T WE 32100.
    | EM_SPARC -- ^ SPARC.
    | EM_386 -- ^ Intel Architecture.
    | EM_68K -- ^ Motorola 68000.
    | EM_88K -- ^ Motorola 88000.
    | EM_860 -- ^ Intel 80860.
    | EM_MIPS -- ^ MIPS RS3000 Big-Endian.
    | EM_MIPS_RS4_BE -- ^ MIPS RS4000 Big-Endian.
    | EM_RISCV -- ^ RISC-V.
    deriving Eq

instance Enum ElfMachine
  where
    toEnum 0x0 = EM_NONE
    toEnum 0x1 = EM_M32
    toEnum 0x2 = EM_SPARC
    toEnum 0x3 = EM_386
    toEnum 0x4 = EM_68K
    toEnum 0x5 = EM_88K
    toEnum 0x7 = EM_860
    toEnum 0x8 = EM_MIPS
    toEnum 0xa = EM_MIPS_RS4_BE
    toEnum 0xf3 = EM_RISCV
    toEnum x = ElfMachine $ fromIntegral x
    fromEnum EM_NONE = 0x0
    fromEnum EM_M32 = 0x1
    fromEnum EM_SPARC = 0x2
    fromEnum EM_386 = 0x3
    fromEnum EM_68K = 0x4
    fromEnum EM_88K = 0x5
    fromEnum EM_860 = 0x7
    fromEnum EM_MIPS = 0x8
    fromEnum EM_MIPS_RS4_BE = 0xa
    fromEnum EM_RISCV = 0xf3
    fromEnum (ElfMachine x) = fromIntegral x

data ElfVersion
    = ElfVersion Elf32_Word
    | EV_NONE -- ^ Invalid versionn.
    | EV_CURRENT -- ^ Current version.
    deriving Eq

instance Enum ElfVersion
  where
    toEnum 0 = EV_NONE
    toEnum 1 = EV_CURRENT
    toEnum x = ElfVersion $ fromIntegral x
    fromEnum EV_NONE = 0
    fromEnum EV_CURRENT = 1
    fromEnum (ElfVersion x) = fromIntegral x

data ElfType
    = ElfType Elf32_Half
    | ET_NONE -- ^ No file type.
    | ET_REL -- ^ Relocatable file.
    | ET_EXEC -- ^ Executable file.
    | ET_DYN -- ^ Shared object file.
    | ET_CORE -- ^ Core file.
    | ET_LOPROC -- ^ Processor-specific.
    | ET_HIPROC -- ^ Processor-specific.
    deriving Eq

instance Enum ElfType
  where
    toEnum 0 = ET_NONE
    toEnum 1 = ET_REL
    toEnum 2 = ET_EXEC
    toEnum 3 = ET_DYN
    toEnum 4 = ET_CORE
    toEnum 0xff00 = ET_LOPROC
    toEnum 0xffff = ET_HIPROC
    toEnum x = ElfType $ fromIntegral x
    fromEnum ET_NONE = 0
    fromEnum ET_REL = 1
    fromEnum ET_EXEC = 2
    fromEnum ET_DYN = 3
    fromEnum ET_CORE = 4
    fromEnum ET_LOPROC = 0xff00
    fromEnum ET_HIPROC = 0xffff
    fromEnum (ElfType x) = fromIntegral x

data Elf32_Sym = Elf32_Sym
    { st_name :: Elf32_Word
    , st_value :: Elf32_Addr
    , st_size :: Elf32_Word
    , st_info :: Word8
    , st_other :: Word8
    , st_shndx :: Elf32_Half
    } deriving Eq

data ElfSymbolBinding
    = ElfSymbolBinding Word8
    | STB_LOCAL
    | STB_GLOBAL
    | STB_WEAK
    | STB_LOPROC
    | STB_HIPROC
    deriving Eq

instance Enum ElfSymbolBinding
  where
    toEnum 0 = STB_LOCAL
    toEnum 1 = STB_GLOBAL
    toEnum 2 = STB_WEAK
    toEnum 13 = STB_LOPROC
    toEnum 15 = STB_HIPROC
    toEnum x = ElfSymbolBinding $ fromIntegral x
    fromEnum STB_LOCAL = 0
    fromEnum STB_GLOBAL = 1
    fromEnum STB_WEAK = 2
    fromEnum STB_LOPROC = 13
    fromEnum STB_HIPROC = 15
    fromEnum (ElfSymbolBinding x) = fromIntegral x

data ElfSymbolType
    = ElfSymbolType Word8
    | STT_NOTYPE
    | STT_OBJECT
    | STT_FUNC
    | STT_SECTION
    | STT_FILE
    | STT_LOPROC
    | STT_HIPROC
    deriving Eq

instance Enum ElfSymbolType
  where
    toEnum 0 = STT_NOTYPE
    toEnum 1 = STT_OBJECT
    toEnum 2 = STT_FUNC
    toEnum 3 = STT_SECTION
    toEnum 4 = STT_FILE
    toEnum 13 = STT_LOPROC
    toEnum 15 = STT_HIPROC
    toEnum x = ElfSymbolType $ fromIntegral x
    fromEnum STT_NOTYPE = 0
    fromEnum STT_OBJECT = 1
    fromEnum STT_FUNC = 2
    fromEnum STT_SECTION = 3
    fromEnum STT_FILE = 4
    fromEnum STT_LOPROC = 13
    fromEnum STT_HIPROC = 15
    fromEnum (ElfSymbolType x) = fromIntegral x

data Elf32_Rel = Elf32_Rel
    { r_offset :: Elf32_Addr
    , r_info :: Elf32_Word
    } deriving Eq

data Elf32_Rela = Elf32_Rela
    { r_offset :: Elf32_Addr
    , r_info :: Elf32_Word
    , r_addend :: Elf32_Sword
    } deriving Eq

data ElfSectionType
    = ElfSectionType Elf32_Word
    | SHT_NULL
    | SHT_PROGBITS
    | SHT_SYMTAB
    | SHT_STRTAB
    | SHT_RELA
    | SHT_HASH
    | SHT_DYNAMIC
    | SHT_NOTE
    | SHT_NOBITS
    | SHT_REL
    | SHT_SHLIB
    | SHT_DYNSYM
    | SHT_LOPROC
    | SHT_HIPROC
    | SHT_LOUSER
    | SHT_HIUSER
    deriving Eq

instance Enum ElfSectionType
  where
    toEnum 0 = SHT_NULL
    toEnum 1 = SHT_PROGBITS
    toEnum 2 = SHT_SYMTAB
    toEnum 3 = SHT_STRTAB
    toEnum 4 = SHT_RELA
    toEnum 5 = SHT_HASH
    toEnum 6 = SHT_DYNAMIC
    toEnum 7 = SHT_NOTE
    toEnum 8 = SHT_NOBITS
    toEnum 9 = SHT_REL
    toEnum 10 = SHT_SHLIB
    toEnum 11 = SHT_DYNSYM
    toEnum 0x70000000 = SHT_LOPROC
    toEnum 0x7fffffff = SHT_HIPROC
    toEnum 0x80000000 = SHT_LOUSER
    toEnum 0xffffffff = SHT_HIUSER
    toEnum x = ElfSectionType $ fromIntegral x
    fromEnum SHT_NULL = 0
    fromEnum SHT_PROGBITS = 1
    fromEnum SHT_SYMTAB = 2
    fromEnum SHT_STRTAB = 3
    fromEnum SHT_RELA = 4
    fromEnum SHT_HASH = 5
    fromEnum SHT_DYNAMIC = 6
    fromEnum SHT_NOTE = 7
    fromEnum SHT_NOBITS = 8
    fromEnum SHT_REL = 9
    fromEnum SHT_SHLIB = 10
    fromEnum SHT_DYNSYM = 11
    fromEnum SHT_LOPROC = 0x70000000
    fromEnum SHT_HIPROC = 0x7fffffff
    fromEnum SHT_LOUSER = 0x80000000
    fromEnum SHT_HIUSER = 0xffffffff
    fromEnum (ElfSectionType x) = fromIntegral x

-- * Constants.

shfWrite :: Elf32_Word
shfWrite = 0x1

shfAlloc :: Elf32_Word
shfAlloc = 0x2

shfExecinstr:: Elf32_Word
shfExecinstr = 0x4

shfMascproc :: Elf32_Word
shfMascproc = 0xf0000000

shfInfoLink :: Elf32_Word
shfInfoLink = 0x40

-- * Encoding functions.

elf32Addr :: ByteOrder -> Elf32_Addr -> ByteString.Builder.Builder
elf32Addr LSB = ByteString.Builder.word32LE
elf32Addr MSB = ByteString.Builder.word32BE

elf32Half :: ByteOrder -> Elf32_Half -> ByteString.Builder.Builder
elf32Half LSB = ByteString.Builder.word16LE
elf32Half MSB = ByteString.Builder.word16BE

elf32Off :: ByteOrder -> Elf32_Off -> ByteString.Builder.Builder
elf32Off LSB = ByteString.Builder.word32LE
elf32Off MSB = ByteString.Builder.word32BE

elf32Sword :: ByteOrder -> Elf32_Sword -> ByteString.Builder.Builder
elf32Sword LSB = ByteString.Builder.int32LE
elf32Sword MSB = ByteString.Builder.int32BE

elf32Word :: ByteOrder -> Elf32_Word -> ByteString.Builder.Builder
elf32Word LSB = ByteString.Builder.word32LE
elf32Word MSB = ByteString.Builder.word32BE

elfIdentification :: ElfIdentification -> ByteString.Builder.Builder
elfIdentification (ElfIdentification elfClass elfData)
    = ByteString.Builder.word8 0x7f
    <> ByteString.Builder.string7 "ELF"
    <> ByteString.Builder.word8 (fromIntegralEnum elfClass)
    <> ByteString.Builder.word8 (fromIntegralEnum elfData)
    <> ByteString.Builder.word8 (fromIntegralEnum EV_CURRENT)
    <> ByteString.Builder.byteString (ByteString.replicate 9 0)

elf32Ehdr :: Elf32_Ehdr -> Either ElfEncodingError ByteString.Builder.Builder
elf32Ehdr Elf32_Ehdr{..} = encode <$> byteOrder e_ident
  where
    encode byteOrder'
        = elfIdentification e_ident
        <> elf32Half byteOrder' (fromIntegralEnum e_type)
        <> elf32Half byteOrder' (fromIntegralEnum e_machine)
        <> elf32Word byteOrder' (fromIntegralEnum e_version)
        <> elf32Addr byteOrder' e_entry
        <> elf32Off byteOrder' e_phoff
        <> elf32Off byteOrder' e_shoff
        <> elf32Word byteOrder' e_flags
        <> elf32Half byteOrder' e_ehsize
        <> elf32Half byteOrder' e_phentsize
        <> elf32Half byteOrder' e_phnum
        <> elf32Half byteOrder' e_shentsize
        <> elf32Half byteOrder' e_shnum
        <> elf32Half byteOrder' e_shstrndx

byteOrder :: ElfIdentification -> Either ElfEncodingError ByteOrder
byteOrder (ElfIdentification class' _)
    | class' /= ELFCLASS32 = Left $ ElfUnsupportedClassError class'
byteOrder (ElfIdentification _ ELFDATA2MSB) = Right MSB
byteOrder (ElfIdentification _ ELFDATA2LSB) = Right LSB
byteOrder (ElfIdentification _ ELFDATANONE) = Left ElfInvalidByteOrderError

elf32Shdr :: ByteOrder -> Elf32_Shdr -> ByteString.Builder.Builder
elf32Shdr byteOrder' Elf32_Shdr{..}
    = elf32Word byteOrder' sh_name
    <> elf32Word byteOrder' (fromIntegralEnum sh_type)
    <> elf32Word byteOrder' sh_flags
    <> elf32Addr byteOrder' sh_addr
    <> elf32Off byteOrder' sh_offset
    <> elf32Word byteOrder' sh_size
    <> elf32Word byteOrder' sh_link
    <> elf32Word byteOrder' sh_info
    <> elf32Word byteOrder' sh_addralign
    <> elf32Word byteOrder' sh_entsize

elf32Sym :: ByteOrder -> Elf32_Sym -> ByteString.Builder.Builder
elf32Sym byteOrder' Elf32_Sym{..}
    = elf32Word byteOrder' st_name
    <> elf32Addr byteOrder' st_value
    <> elf32Word byteOrder' st_size
    <> ByteString.Builder.word8 st_info
    <> ByteString.Builder.word8 st_other
    <> elf32Half byteOrder' st_shndx

elf32Rel :: ByteOrder -> Elf32_Rel -> ByteString.Builder.Builder
elf32Rel byteOrder' Elf32_Rel{..}
    = elf32Addr byteOrder' r_offset
    <> elf32Word byteOrder' r_info

elf32Rela :: ByteOrder -> Elf32_Rela -> ByteString.Builder.Builder
elf32Rela byteOrder' Elf32_Rela{..}
    = elf32Addr byteOrder' r_offset
    <> elf32Word byteOrder' r_info
    <> elf32Sword byteOrder' r_addend

stInfo :: ElfSymbolBinding -> ElfSymbolType -> Word8
stInfo binding type' = fromIntegralEnum binding `shiftL` 4
    .|. (fromIntegralEnum type' .&. 0xf)

rInfo :: Elf32_Word -> Word8 -> Elf32_Word
rInfo symbol type' = symbol `shiftL` 8
    .|. fromIntegralEnum type'

-- * Help types and functions.

data ByteOrder = LSB | MSB
    deriving Eq

data ElfEncodingError
    = ElfInvalidByteOrderError
    | ElfUnsupportedClassError ElfClass
    deriving Eq

instance Show ElfEncodingError
  where
    show ElfInvalidByteOrderError = "Invalid byte order."
    show (ElfUnsupportedClassError class') =
        concat ["Elf class \"", show class', "\" is not supported."]

instance Exception ElfEncodingError

fromIntegralEnum :: (Enum a, Num b) => a -> b
fromIntegralEnum = fromIntegral . fromEnum

-- * Object file generation.

newtype ElfWriter a = ElfWriter
    { runElfWriter :: StateT (ElfHeaderResult Elf32_Shdr) IO a
    }

data ElfHeaderResult a = ElfHeaderResult
    { sectionNames :: ByteString
    , 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 :: ByteString -> Elf32_Shdr -> ElfWriter ()
addSectionHeader name newHeader = ElfWriter $ modify' modifier
  where
    modifier ElfHeaderResult{..} =
        ElfHeaderResult
            { sectionHeaders = Vector.snoc sectionHeaders newHeader
            , sectionNames = sectionNames <> name <> "\0"
            }

-- 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 -> (Handle -> ElfWriter Elf32_Ehdr) -> IO ()
elfObject outFile putContents = withFile outFile WriteMode withObjectFile
  where
    withObjectFile objectHandle
        = hSeek objectHandle AbsoluteSeek (fromIntegral elfHeaderSize)
        >> putContents' objectHandle
        >>= uncurry (afterContents objectHandle)
    putContents' objectHandle
        = flip runStateT initialState
        $ runElfWriter
        $ putContents objectHandle
    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 = ElfHeaderResult
        { sectionHeaders = Vector.singleton zeroHeader
        , sectionNames = "\0"
        }
    afterContents objectHandle header 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