{-# LANGUAGE BlockArguments #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE TypeFamilies #-}

module Main (main) where

-- import Paths_elf
import Prelude as P

import Control.Monad
import Control.Monad.Catch
import Data.Binary
import qualified Data.ByteString as BS
import Data.ByteString.Lazy as BSL
-- import Data.ByteString.Lazy.Char8 as BSC
import Data.Foldable as F
-- import Data.Functor.Identity
import Data.Int
import Data.Singletons
import Data.Singletons.Sigma
import Prettyprinter
import Prettyprinter.Render.Text
import System.Directory
import System.FilePath
import System.IO as IO
import Test.Tasty
import Test.Tasty.Golden
import Test.Tasty.HUnit

import Control.Exception.ChainedException
import Data.Elf
import Data.Elf.PrettyPrint
import Data.Elf.Headers
import Data.Endian

partitionM :: Monad m => (a -> m Bool) -> [a] -> m ([a], [a])
partitionM p = foldlM f ([], [])
    where
        f (ts, fs) x = do
            b <- p x
            return $ if b then (x:ts, fs) else (ts, x:fs)

traverseDir :: FilePath -> (FilePath -> IO Bool) -> IO [FilePath]
traverseDir root ok = go root
    where
        go :: FilePath -> IO [FilePath]
        go dir = do
            paths <- P.map (dir </>) <$> listDirectory dir
            (dirPaths, filePaths) <- partitionM doesDirectoryExist paths
            oks <- filterM ok filePaths
            (oks ++) <$> (F.concat <$> mapM go dirPaths)

isElf :: FilePath -> IO Bool
isElf p = if takeExtension p == ".bad" then return False else (elfMagic ==) . decode <$> BSL.readFile p

decodeOrFailAssertion :: Binary a => ByteString -> IO (Int64, a)
decodeOrFailAssertion bs = case decodeOrFail bs of
    Left (_, off, err) -> assertFailure (err ++ " @" ++ show off)
    Right (_, off, a) -> return (off, a)

mkTest'' :: forall a . IsElfClass a => HeaderXX a -> ByteString -> Assertion
mkTest'' HeaderXX{..} bs = do

    let
        takeLen off len = BSL.take (fromIntegral len) $ BSL.drop (fromIntegral off) bs
        bsSections = takeLen hShOff (hShEntSize * hShNum)
        bsSegments = takeLen hPhOff (hPhEntSize * hPhNum)

    (ss :: [SectionXX a]) <- parseBList hData bsSections
    assertEqual "Section table round trip does not work" bsSections $ serializeBList hData ss

    (ps :: [SegmentXX a]) <- parseBList hData bsSegments
    assertEqual "Segment table round trip does not work" bsSegments $ serializeBList hData ps

mkTest' :: ByteString -> Assertion
mkTest' bs = do
    (off, elfh@(classS :&: hxx) :: Header) <- decodeOrFailAssertion bs
    assertBool "Incorrect header size" (headerSize (fromSing classS) == off)
    assertEqual "Header round trip does not work" (BSL.take off bs) (encode elfh)

    withElfClass classS mkTest'' hxx bs

mkTest :: FilePath -> TestTree
mkTest p = testCase p $ withBinaryFile p ReadMode (BSL.hGetContents >=> mkTest')

mkGoldenTest' :: FilePath -> FilePath -> (FilePath -> IO (Doc ())) -> FilePath -> TestTree
mkGoldenTest' g o formatFunction file = goldenVsFile file g o mkGoldenTestOutput
    where
        mkGoldenTestOutput :: IO ()
        mkGoldenTestOutput = do
            doc <- formatFunction file
            withFile o WriteMode (`hPutDoc` doc)

mkGoldenTest :: String -> (FilePath -> IO (Doc ())) -> FilePath -> TestTree
mkGoldenTest name formatFunction file = mkGoldenTest' g o formatFunction file
    where
        newBase = "tests" </> file <.> name
        o = newBase <.> "out"
        g = newBase <.> "golden"

mkGoldenTestOSuffix :: String -> String -> (FilePath -> IO (Doc ())) -> FilePath -> TestTree
mkGoldenTestOSuffix name osuffix formatFunction file = mkGoldenTest' g o formatFunction file
    where
        newBase = "tests" </> file <.> name
        o = newBase <.> osuffix <.> "out"
        g = newBase <.> "golden"

------------------------------------------------------------------------------

-- FIXME: define foldMapRBuilderList

index' :: (Integral i, MonadThrow m) => [a] -> i -> m a
index' (x:_) 0 = return x
index' (_:xs) n | n > 0     = index' xs (n-1)
                | otherwise = $chainedError "index': negative argument."
index' _ _                  = $chainedError "index': index too large."

getStringTable :: MonadThrow m => Sigma ElfClass (TyCon1 HeadersXX) -> BSL.ByteString -> m BSL.ByteString
getStringTable (classS :&: HeadersXX (HeaderXX{..}, ss, _)) bs = withElfClass classS
    if hShStrNdx == 0
        then return BSL.empty
        else do
            strs <- index' ss hShStrNdx
            return $ getSectionData bs strs

copyElf :: MonadCatch m => BSL.ByteString -> m BSL.ByteString
copyElf bs = parseElf bs >>= serializeElf

---------------------------------------------------------------------

-- This is for examples/README.md
withHeader :: BSL.ByteString -> (forall a . IsElfClass a => HeaderXX a -> b) -> Either String b
withHeader bs f =
    case decodeOrFail bs of
        Left (_, _, err) -> Left err
        Right (_, _, (classS :&: hxx) :: Header) -> Right $ withElfClass classS f hxx

printHeaderFile :: FilePath -> IO (Doc ())
printHeaderFile path = do
    bs <- fromStrict <$> BS.readFile path
    $eitherAddContext' $ withHeader bs printHeader

printHeadersFile :: FilePath -> IO (Doc ())
printHeadersFile path = do
    bs <- fromStrict <$> BS.readFile path
    (classS :&: HeadersXX (hdr, ss, ps)) <- parseHeaders bs
    return $ withSingI classS $ printHeaders hdr ss ps

printStrTableFile :: FilePath -> IO (Doc ())
printStrTableFile path = do
    bs <- readFileLazy path
    hdrs <- parseHeaders bs
    st <- getStringTable hdrs bs
    printStringTable st

printCopyStrTableFile :: FilePath -> IO (Doc ())
printCopyStrTableFile path = do
    bs <- readFileLazy path
    bs' <- copyElf bs
    hdrs <- parseHeaders bs'
    st <- getStringTable hdrs bs'
    printStringTable st

printRBuilderFile :: FilePath -> IO (Doc ())
printRBuilderFile path = do
    bs <- readFileLazy path
    hdrs <- parseHeaders bs
    printLayout hdrs bs

printCopyRBuilderFile :: FilePath -> IO (Doc ())
printCopyRBuilderFile path = do
    bs <- readFileLazy path
    bs' <- copyElf bs
    hdrs <- parseHeaders bs'
    printLayout hdrs bs'

printElfFile :: FilePath -> IO (Doc ())
printElfFile path = do
    bs <- readFileLazy path
    e <- parseElf bs
    printElf e

printCopyElfFile :: FilePath -> IO (Doc ())
printCopyElfFile path = do
    bs <- readFileLazy path
    bs' <- copyElf bs
    e <- parseElf bs'
    printElf e

-----------------------------------------------------------------------

testHeader64 :: Header
testHeader64 = SELFCLASS64 :&: HeaderXX ELFDATA2LSB 0 0 0 0 0 0 0 0 0 0 0 0 0

testHeader32 :: Header
testHeader32 = SELFCLASS32 :&: HeaderXX ELFDATA2MSB 0 0 0 0 0 0 0 0 0 0 0 0 0

testSection64 :: SectionXX 'ELFCLASS64
testSection64 = SectionXX 0 0 0 0 0 0 0 0 0 0

testSection32 :: SectionXX 'ELFCLASS32
testSection32 = SectionXX 0 0 0 0 0 0 0 0 0 0

testSegment64 :: SegmentXX 'ELFCLASS64
testSegment64 =  SegmentXX 0 0 0 0 0 0 0 0

testSegment32 :: SegmentXX 'ELFCLASS32
testSegment32 =  SegmentXX 0 0 0 0 0 0 0 0

testSymbolTableEntry64 :: SymbolXX 'ELFCLASS64
testSymbolTableEntry64 =  SymbolXX 0 0 0 0 0 0

testSymbolTableEntry32 :: SymbolXX 'ELFCLASS32
testSymbolTableEntry32 =  SymbolXX 0 0 0 0 0 0

mkSizeTest :: Binary a => String -> a -> Int64 -> TestTree
mkSizeTest name v s = testCase name (len @?= s)
    where
        len = BSL.length $ encode v

hdrSizeTests :: TestTree
hdrSizeTests = testGroup "header size" [ mkSizeTest "header 64" testHeader64 (headerSize ELFCLASS64)
                                       , mkSizeTest "header 32" testHeader32 (headerSize ELFCLASS32)

                                       , mkSizeTest "section 64" (Le testSection64) (sectionTableEntrySize ELFCLASS64)
                                       , mkSizeTest "section 32" (Be testSection32) (sectionTableEntrySize ELFCLASS32)

                                       , mkSizeTest "segment 64" (Le testSegment64) (segmentTableEntrySize ELFCLASS64)
                                       , mkSizeTest "segment 32" (Be testSegment32) (segmentTableEntrySize ELFCLASS32)

                                       , mkSizeTest "symbol table entry 64" (Le testSymbolTableEntry64) (symbolTableEntrySize ELFCLASS64)
                                       , mkSizeTest "symbol table entry 32" (Be testSymbolTableEntry32) (symbolTableEntrySize ELFCLASS32)
                                       ]

elfsForHeader :: [String]
elfsForHeader = [ "testdata/orig/bloated"
                , "testdata/orig/tiny"
                , "testdata/orig/vdso"
                ]

main :: IO ()
main = do

    elfs <- traverseDir "testdata" isElf

    defaultMain $ testGroup "elf" [ hdrSizeTests
                                  , testGroup "headers round trip"  (mkTest <$> elfs)
                                  , testGroup "elf headers golden"  (mkGoldenTest        "elf_header"      printHeaderFile       <$> elfsForHeader)
                                  , testGroup "header golden"       (mkGoldenTest        "header"          printHeadersFile      <$> elfs)
                                  , testGroup "string table golden" (mkGoldenTest        "strtable"        printStrTableFile     <$> elfs)
                                  , testGroup "layout golden"       (mkGoldenTest        "layout"          printRBuilderFile     <$> elfs)
                                  , testGroup "elf golden"          (mkGoldenTest        "elf"             printElfFile          <$> elfs)
                                  , testGroup "string table copy"   (mkGoldenTestOSuffix "strtable" "copy" printCopyStrTableFile <$> elfs)
                                  , testGroup "layout copy"         (mkGoldenTestOSuffix "layout"   "copy" printCopyRBuilderFile <$> elfs)
                                  , testGroup "elf copy"            (mkGoldenTestOSuffix "elf"      "copy" printCopyElfFile      <$> elfs)
                                  ]