{-# LANGUAGE PackageImports #-}

module Main (main) where

-- TODO: bench bzlib?
import qualified "bz2" Codec.Compression.BZip   as BZ2
#ifdef WITH_BZLIB
import qualified "bzlib" Codec.Compression.BZip as BZlib
#endif
import           Conduit                        (runResourceT)
import           Criterion.Main
import qualified Data.ByteString.Lazy           as BSL
import           Data.Conduit                   (runConduit, (.|))
import qualified Data.Conduit.BZlib             as C
import           Data.Conduit.Combinators       (sinkFile, sourceFile)
import           Pipes                          (runEffect, (>->))
import qualified Pipes.ByteString               as P
import qualified Pipes.BZip                     as P
import           Pipes.Safe                     (runSafeT)
import           System.FilePath                ((</>))
import           System.IO                      (IOMode (ReadMode, WriteMode), withFile)
import           System.IO.Temp                 (withSystemTempDirectory)

decompressDump :: (BSL.ByteString -> BSL.ByteString) -> IO ()
decompressDump go = withSystemTempDirectory "bz2" $
    \fp -> BSL.writeFile (fp </> "valgrind-3.15.0.tar") =<<
        (go <$> BSL.readFile "valgrind-3.15.0.tar.bz2")

compressDump :: (BSL.ByteString -> BSL.ByteString) -> IO ()
compressDump go = withSystemTempDirectory "bz2" $
    \fp -> BSL.writeFile (fp </> "valgrind-3.15.0.tar.bz2") =<<
        (go <$> BSL.readFile "valgrind-3.15.0.tar")

decompressDumpPipes :: IO ()
decompressDumpPipes = withSystemTempDirectory "bz2" $ \fp ->
    withFile "valgrind-3.15.0.tar.bz2" ReadMode $ \hIn ->
        withFile (fp </> "valgrind-3.15.0.tar") WriteMode $ \hOut ->
            runSafeT $ runEffect $ P.bunzip2 (P.fromHandle hIn) >-> P.toHandle hOut

decompressDumpConduit :: IO ()
decompressDumpConduit = withSystemTempDirectory "bz2" $ \fp ->
    runResourceT $ runConduit $ sourceFile "valgrind-3.15.0.tar.bz2" .| C.bunzip2 .| sinkFile (fp </> "valgrind-3.15.0.tar")

decompressFile :: FilePath -> IO BSL.ByteString
decompressFile = fmap BZ2.decompress . BSL.readFile

compressFile :: FilePath -> IO BSL.ByteString
compressFile = fmap BZ2.compress . BSL.readFile

main :: IO ()
main =
    defaultMain [ bgroup "decompressDump"
                      [ bench "decompress + write to file (bz2)" $ nfIO (decompressDump BZ2.decompress)
#ifdef WITH_BZLIB
                      , bench "decompress + write to file (bzlib)" $ nfIO (decompressDump BZlib.decompress)
#endif
                      , bench "decompress + write to file (pipes-bzip)" $ nfIO decompressDumpPipes
                      , bench "decompress + write to file (bzlib-conduit)" $ nfIO decompressDumpConduit
                      ]
                , bgroup "compressDump"
                      [ bench "compress + write to file (bz2)" $ nfIO (compressDump BZ2.compress)
#ifdef WITH_BZLIB
                      , bench "compress + write to file (bzlib)" $ nfIO (compressDump BZlib.compress)
#endif
                      ]
                , bgroup "decompress"
                      [ bench "decompress file" $ nfIO (decompressFile "test/data/sample1.bz2")
                      , bench "decompress file" $ nfIO (decompressFile "test/data/sample2.bz2")
                      , bench "decompress file" $ nfIO (decompressFile "test/data/sample3.bz2")
                      ]
                , bgroup "compress"
                      [ bench "compress file" $ nfIO (compressFile "test/data/sample1.ref")
                      , bench "compress file" $ nfIO (compressFile "test/data/sample2.ref")
                      , bench "compress file" $ nfIO (compressFile "test/data/sample3.ref")
                      ]
                ]