module Codec.Compression.BZip.Unpack ( decompress
                                     ) where

import Codec.Compression.BZip.Foreign
import Codec.Compression.BZip.Common
import Control.Applicative
import Control.Arrow (second)
import Control.Exception (bracket)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import qualified Data.ByteString.Unsafe as BS
import Foreign.Ptr (Ptr, castPtr, nullPtr)
import Foreign.Marshal.Alloc (free, mallocBytes)
import System.IO.Unsafe (unsafeDupablePerformIO)

#include <bzlib.h>

{-# NOINLINE decompress #-}
-- | Don't use this on pathological input; it may not be secure
--
-- @since 0.1.1.0
decompress :: BSL.ByteString -> BSL.ByteString
decompress bsl = unsafeDupablePerformIO $
    let bss = BSL.toChunks bsl in
    BSL.fromChunks <$> bracket
        (do { p <- bzStreamInit ; bzDecompressInit p ; pure p })
        bZ2BzDecompressEnd
        (\p -> bzDecompressChunks p bss)

bzDecompressChunks :: Ptr BzStream -> [BS.ByteString] -> IO [BS.ByteString]
bzDecompressChunks p bs =

    bracket
        (mallocBytes bufSz)
        free
        (fmap snd . extractBuf bs)

    where

        -- corresponds to inner loop in zlib example
        fillBuf :: [BS.ByteString] -> Ptr a -> IO (BZError, [BS.ByteString])
        fillBuf bs' bufOut = do

            {# set bz_stream.avail_out #} p bufSz
            {# set bz_stream.next_out #} p (castPtr bufOut)

            ret <- bZ2BzDecompress p

            szOut <- fromIntegral <$> {# get bz_stream->avail_out #} p

            let bytesAvail = bufSz - szOut

            newBSAp <- if bytesAvail /= 0
                then (:) <$> BS.packCStringLen (castPtr bufOut, bytesAvail)
                else pure id

            if ret == BzStreamEnd
                then pure (ret, newBSAp [])
                else if szOut == 0
                    then second newBSAp <$> fillBuf bs' bufOut
                    else second newBSAp <$> extractBuf bs' bufOut

        -- corresponds to outer loop in zlib example
        extractBuf :: [BS.ByteString] -> Ptr a -> IO (BZError, [BS.ByteString])
        extractBuf [] bufOut = do

            (res, blocks) <- fillBuf [] bufOut

            if res == BzStreamEnd
                then pure (BzStreamEnd, blocks)
                else extractBuf [] bufOut

        extractBuf (b:bs') bufOut =
            BS.unsafeUseAsCStringLen b $ \(buf, sz) -> do

                {# set bz_stream.avail_in #} p (fromIntegral sz)
                {# set bz_stream.next_in #} p buf

                (res, blocks) <- fillBuf bs' bufOut

                if res == BzStreamEnd
                    then pure (BzStreamEnd, blocks)
                    else extractBuf bs' bufOut

        bufSz :: Integral a => a
        bufSz = 32 * 1024

bzDecompressInit :: Ptr BzStream -> IO ()
bzDecompressInit p = do

    {# set bz_stream.next_in #} p nullPtr
    {# set bz_stream.avail_in #} p 0

    bZ2BzDecompressInit p 0 False