    {-# LANGUAGE FlexibleContexts #-}
    {-# LANGUAGE RankNTypes #-}
    -- | Streaming compression and decompression using conduits.
    --
    -- Parts of this code were taken from zlib-enum and adapted for conduits.
    module Data.Conduit.Zlib (
        -- * Conduits
        compress, decompress, gzip, ungzip,
        -- * Flushing
        compressFlush, decompressFlush,
        -- * Decompression combinators
        multiple,
        -- * Re-exported from zlib-bindings
        WindowBits (..), defaultWindowBits
    ) where

    import Data.Streaming.Zlib
    import Data.Conduit
    import Data.ByteString (ByteString)
    import qualified Data.ByteString as S
    import Control.Monad (unless, liftM)
    import Control.Monad.Trans.Class (lift, MonadTrans)
    import Control.Monad.Primitive (PrimMonad, unsafePrimToPrim)
    import Control.Monad.Base (MonadBase, liftBase)
    import Control.Monad.Trans.Resource (MonadThrow, monadThrow)
    import Data.Function (fix)

    -- | Gzip compression with default parameters.
    gzip :: (MonadThrow m, MonadBase base m, PrimMonad base) => Conduit ByteString m ByteString
    gzip = compress 1 (WindowBits 31)

    -- | Gzip decompression with default parameters.
    ungzip :: (MonadBase base m, PrimMonad base, MonadThrow m) => Conduit ByteString m ByteString
    ungzip = decompress (WindowBits 31)

    unsafeLiftIO :: (MonadBase base m, PrimMonad base, MonadThrow m) => IO a -> m a
    unsafeLiftIO = liftBase . unsafePrimToPrim

    -- |
    -- Decompress (inflate) a stream of 'ByteString's. For example:
    --
    -- >    sourceFile "test.z" $= decompress defaultWindowBits $$ sinkFile "test"

    decompress
        :: (MonadBase base m, PrimMonad base, MonadThrow m)
        => WindowBits -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)
        -> Conduit ByteString m ByteString
    decompress =
        helperDecompress (liftM (fmap Chunk) await) yield' leftover
      where
        yield' Flush = return ()
        yield' (Chunk bs) = yield bs

    -- | Same as 'decompress', but allows you to explicitly flush the stream.
    decompressFlush
        :: (MonadBase base m, PrimMonad base, MonadThrow m)
        => WindowBits -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)
        -> Conduit (Flush ByteString) m (Flush ByteString)
    decompressFlush = helperDecompress await yield (leftover . Chunk)

    helperDecompress :: (Monad (t m), MonadBase base m, PrimMonad base, MonadThrow m, MonadTrans t)
                     => t m (Maybe (Flush ByteString))
                     -> (Flush ByteString -> t m ())
                     -> (ByteString -> t m ())
                     -> WindowBits
                     -> t m ()
    helperDecompress await' yield' leftover' config = do
        -- Initialize the stateful inflater, which will be used below
        -- This inflater is never exposed outside of this function
        inf <- lift $ unsafeLiftIO $ initInflate config

        -- Some helper functions used by the main feeder loop below

        let -- Flush any remaining inflated bytes downstream
            flush = do
                chunk <- lift $ unsafeLiftIO $ flushInflate inf
                unless (S.null chunk) $ yield' $ Chunk chunk

            -- Get any input which is unused by the inflater
            getUnused = lift $ unsafeLiftIO $ getUnusedInflate inf

            -- If there is any unused data, return it as leftovers to the stream
            unused = do
                rem' <- getUnused
                unless (S.null rem') $ leftover' rem'

        -- Main loop: feed data from upstream into the inflater
        fix $ \feeder -> do
            mnext <- await'
            case mnext of
                -- No more data is available from upstream
                Nothing -> do
                    -- Flush any remaining uncompressed data
                    flush
                    -- Return the rest of the unconsumed data as leftovers
                    unused
                -- Another chunk of compressed data arrived
                Just (Chunk x) -> do
                    -- Feed the compressed data into the inflater, returning a
                    -- "popper" which will return chunks of decompressed data
                    popper <- lift $ unsafeLiftIO $ feedInflate inf x

                    -- Loop over the popper grabbing decompressed chunks and
                    -- yielding them downstream
                    fix $ \pop -> do
                        mbs <- lift $ unsafeLiftIO popper
                        case mbs of
                            -- No more data from this popper
                            PRDone -> do
                                rem' <- getUnused
                                if S.null rem'
                                    -- No data was unused by the inflater, so let's
                                    -- fill it up again and get more data out of it
                                    then feeder
                                    -- In this case, there is some unconsumed data,
                                    -- meaning the compressed stream is complete.
                                    -- At this point, we need to stop feeding,
                                    -- return the unconsumed data as leftovers, and
                                    -- flush any remaining content (which should be
                                    -- nothing)
                                    else do
                                        flush
                                        leftover' rem'
                            -- Another chunk available, yield it downstream and
                            -- loop again
                            PRNext bs -> do
                                yield' (Chunk bs)
                                pop
                            -- An error occurred inside zlib, throw it
                            PRError e -> lift $ monadThrow e
                -- We've been asked to flush the stream
                Just Flush -> do
                    -- Get any uncompressed data waiting for us
                    flush
                    -- Put a Flush in the stream
                    yield' Flush
                    -- Feed in more data
                    feeder

    -- |
    -- Compress (deflate) a stream of 'ByteString's. The 'WindowBits' also control
    -- the format (zlib vs. gzip).

    compress
        :: (MonadBase base m, PrimMonad base, MonadThrow m)
        => Int         -- ^ Compression level
        -> WindowBits  -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)
        -> Conduit ByteString m ByteString
    compress =
        helperCompress (liftM (fmap Chunk) await) yield'
      where
        yield' Flush = return ()
        yield' (Chunk bs) = yield bs

    -- | Same as 'compress', but allows you to explicitly flush the stream.
    compressFlush
        :: (MonadBase base m, PrimMonad base, MonadThrow m)
        => Int         -- ^ Compression level
        -> WindowBits  -- ^ Zlib parameter (see the zlib-bindings package as well as the zlib C library)
        -> Conduit (Flush ByteString) m (Flush ByteString)
    compressFlush = helperCompress await yield

    helperCompress :: (Monad (t m), MonadBase base m, PrimMonad base, MonadThrow m, MonadTrans t)
                   => t m (Maybe (Flush ByteString))
                   -> (Flush ByteString -> t m ())
                   -> Int
                   -> WindowBits
                   -> t m ()
    helperCompress await' yield' level config =
        await' >>= maybe (return ()) start
      where
        start input = do
            def <- lift $ unsafeLiftIO $ initDeflate level config
            push def input

        continue def = await' >>= maybe (close def) (push def)

        goPopper popper = do
            mbs <- lift $ unsafeLiftIO popper
            case mbs of
                PRDone -> return ()
                PRNext bs -> yield' (Chunk bs) >> goPopper popper
                PRError e -> lift $ monadThrow e

        push def (Chunk x) = do
            popper <- lift $ unsafeLiftIO $ feedDeflate def x
            goPopper popper
            continue def

        push def Flush = do
            mchunk <- lift $ unsafeLiftIO $ flushDeflate def
            case mchunk of
                PRDone -> return ()
                PRNext x -> yield' $ Chunk x
                PRError e -> lift $ monadThrow e
            yield' Flush
            continue def

        close def = do
            mchunk <- lift $ unsafeLiftIO $ finishDeflate def
            case mchunk of
                PRDone -> return ()
                PRNext chunk -> yield' (Chunk chunk) >> close def
                PRError e -> lift $ monadThrow e

    -- | The standard 'decompress' and 'ungzip' functions will only decompress a
    -- single compressed entity from the stream. This combinator will exhaust the
    -- stream completely of all individual compressed entities. This is useful for
    -- cases where you have a concatenated archive, e.g. @cat file1.gz file2.gz >
    -- combined.gz@.
    --
    -- Usage:
    --
    -- > sourceFile "combined.gz" $$ multiple ungzip =$ consume
    --
    -- This combinator will not fail on an empty stream. If you want to ensure that
    -- at least one compressed entity in the stream exists, consider a usage such
    -- as:
    --
    -- > sourceFile "combined.gz" $$ (ungzip >> multiple ungzip) =$ consume
    --
    -- @since 1.1.10
    multiple :: Monad m
             => Conduit ByteString m a
             -> Conduit ByteString m a
    multiple inner =
        loop
      where
        loop = do
            mbs <- await
            case mbs of
                Nothing -> return ()
                Just bs
                    | S.null bs -> loop
                    | otherwise -> do
                        leftover bs
                        inner
                        loop