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(..)
    , elf32Addr
    , elf32Half
    , elf32Off
    , elf32Shdr
    , elf32Sword
    , elf32Word
    , elf32Ehdr
    , elf32Rel
    , elf32Rela
    , elf32Sym
    , elfIdentification
    , rInfo
    , stInfo
    ) where

import Control.Exception (Exception(..))
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

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

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

-- * 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
  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 class' _ <- e_ident
        , class' /= ELFCLASS32 = Left $ ElfUnsupportedClassError class'
        | ElfIdentification _ ELFDATA2MSB <- e_ident = Right MSB
        | ElfIdentification _ ELFDATA2LSB <- e_ident = Right LSB
        | ElfIdentification _ ELFDATANONE <- e_ident = 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'