{-# LINE 1 "Data/Iteratee/BZip.hsc" #-}
{-# LANGUAGE CPP #-}
{-# LINE 2 "Data/Iteratee/BZip.hsc" #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE ForeignFunctionInterface #-}

{-# OPTIONS -Wall -fno-warn-unused-do-bind #-}
module Data.Iteratee.BZip
  (
    -- * Enumeratees
    enumCompress,
    enumDecompress,
    -- * Exceptions
    BZipParamsException(..),
    BZipException(..),
    -- * Parameters
    CompressParams(..),
    defaultCompressParams,
    DecompressParams(..),
    defaultDecompressParams,
    BlockSize(..),
    WorkFactor(..)
  )
where

{-# LINE 24 "Data/Iteratee/BZip.hsc" #-}

import Control.Exception
import Control.Monad.Trans
import Data.ByteString as BS
import Data.ByteString.Internal
import Data.Iteratee
import Data.Typeable
import Foreign
import Foreign.C

{-# LINE 38 "Data/Iteratee/BZip.hsc" #-}

-- | Denotes error is user-supplied parameter
data BZipParamsException
    = IncorrectBlockSize !Int
    -- ^ Incorrect block size was chosen
    | IncorrectWorkFactor !Int
    -- ^ Incorrect work size was chosen
    | IncorrectBufferSize !Int
    -- ^ Incorrect buffer size was chosen
    deriving (Eq,Typeable)

-- | Denotes error in compression and decompression
data BZipException
    = ConfigError
    -- ^ bzip2 library internal error
    | MemError
    -- ^ Memory allocation failed
    | DataError
    -- ^ Corrupted input
    | DataErrorMagic
    -- ^ Incorrect magic number
    | Unexpected !Int
    -- ^ Unknown or unexpected error
    | IncorrectState
    -- ^ Incorrect state - denotes error in library
    deriving (Eq,Typeable)

-- | Denotes the flush that can be sent to stream
data BZipFlush = BZipFlush
    deriving (Eq,Typeable)

instance Show BZipFlush where
    show BZipFlush = "bzlib: flush requested"

instance Exception BZipFlush

fromFlush :: BZipFlush -> CInt
fromFlush BZipFlush = 1
{-# LINE 76 "Data/Iteratee/BZip.hsc" #-}

instance Show BZipParamsException where
    show (IncorrectBlockSize size)
        = "bzlib: incorrect block size " ++ show size
    show (IncorrectWorkFactor wf)
        = "bzlib: incorrect work factor " ++ show wf
    show (IncorrectBufferSize size)
        = "bzlib: incorrect buffer size " ++ show size

instance Show BZipException where
    show ConfigError = "bzlib: library is not configure properly"
    show MemError = "bzlib: memory allocation failed"
    show DataError = "bzlib: input is corrupted"
    show DataErrorMagic = "bzlib: magic number does not match"
    show (Unexpected n) = "bzlib: unexpected error " ++ show n
    show IncorrectState = "bzlib: incorrect state"


instance Exception BZipParamsException
instance Exception BZipException 

newtype BZStream = BZStream (ForeignPtr BZStream)
withBZStream :: BZStream -> (Ptr BZStream -> IO a) -> IO a
withBZStream (BZStream fptr) = withForeignPtr fptr

-- Following code is copied from Duncan Coutts bzlib haskell library version
-- 0.5.2.0 ((c) 2006-2008 Duncan Coutts, published on BSD licence) and adapted

-- | Set of parameters for compression. For sane defaults use
-- 'defaultCompressParams'
data CompressParams = CompressParams {
      compressBlockSize :: BlockSize,
      compressWorkFactor :: WorkFactor,
      -- | The size of output buffer. That is the size of 'Chunk's that will be
      -- emitted to inner iterator (except the last 'Chunk').
      compressBufferSize :: !Int
    }

defaultCompressParams :: CompressParams
defaultCompressParams
    = CompressParams DefaultBlockSize DefaultWorkFactor (8*1024)

-- | Set of parameters for decompression. For sane defaults see 
-- 'defaultDecompressParams'.
data DecompressParams = DecompressParams {
      decompressSaveMemory :: !Bool,
      -- | The size of output buffer. That is the size of 'Chunk's that will be
      -- emitted to inner iterator (except the last 'Chunk').
      decompressBufferSize :: !Int
    }

defaultDecompressParams :: DecompressParams
defaultDecompressParams = DecompressParams False (8*1024)

-- | The compression level specify the tradeoff between speed and compression.
data BlockSize
    = DefaultBlockSize
    -- ^ Default compression level set at 6
    | BestSpeed
    -- ^ The fastest compression method (however less compression)
    | BestCompression
    -- ^ The best compression method (however slowest)
    | CompressionLevel !Int
    -- ^ Compression level set by number from 1 to 9

data WorkFactor
    = DefaultWorkFactor
      -- ^ Default work factor (set at 30)
    | WorkFactor !Int
      -- ^ Hand-tuned work factor

fromBlockSize :: BlockSize -> Either BZipParamsException CInt
fromBlockSize DefaultBlockSize = Right $! 6
fromBlockSize BestSpeed = Right $! 1
fromBlockSize BestCompression = Right $! 9
fromBlockSize (CompressionLevel lvl)
    | lvl < 0 || lvl > 250 = Left $! IncorrectBlockSize $! fromIntegral lvl
    | otherwise = Right $! fromIntegral lvl

fromWorkFactor :: WorkFactor -> Either BZipParamsException CInt
fromWorkFactor DefaultWorkFactor = Right $! 0
fromWorkFactor (WorkFactor wf)
    | wf < 0 || wf > 250 = Left $! IncorrectWorkFactor $! fromIntegral wf
    | otherwise =  Right $! fromIntegral wf

fromErrno :: CInt -> Either BZipException Bool
fromErrno (0) = Right $! True
{-# LINE 163 "Data/Iteratee/BZip.hsc" #-}
fromErrno (1) = Right $! True
{-# LINE 164 "Data/Iteratee/BZip.hsc" #-}
fromErrno (2) = Right $! True
{-# LINE 165 "Data/Iteratee/BZip.hsc" #-}
fromErrno (3) = Right $! True
{-# LINE 166 "Data/Iteratee/BZip.hsc" #-}
fromErrno (4) = Right $! False
{-# LINE 167 "Data/Iteratee/BZip.hsc" #-}
fromErrno (-9) = Left $! ConfigError
{-# LINE 168 "Data/Iteratee/BZip.hsc" #-}
fromErrno (-3) = Left $! MemError
{-# LINE 169 "Data/Iteratee/BZip.hsc" #-}
fromErrno (-4) = Left $! DataError
{-# LINE 170 "Data/Iteratee/BZip.hsc" #-}
fromErrno (-5) = Left $! DataErrorMagic
{-# LINE 171 "Data/Iteratee/BZip.hsc" #-}
fromErrno n = Left $! Unexpected $! fromIntegral n

--
-- In following code we go through 7 states. Some of the operations are
-- 'deterministic' like 'insertOut' and some of them depends on input ('fill')
-- or library call.
--
--                                                  (Finished)
--                                                     ^
--                                                     |
--                                                     |
--                                                     | finish
--                                                     |
--              insertOut                fill[1]       |
---  (Initial) -------------> (EmptyIn) -----------> (Finishing)
--         ^                    ^ | ^ |
--         |             run[2] | | | \------------------\
--         |                    | | |                    |
--         |                    | | \------------------\ |
--         |    run[1]          | |        flush[0]    | |
--         \------------------\ | | fill[0]            | | fill[3]
--                            | | |                    | |
--                            | | |                    | |
--               swapOut      | | v       flush[1]     | v
--  (FullOut) -------------> (Invalid) <----------- (Flushing)
--
-- Initial: Initial state, both buffers are empty
-- EmptyIn: Empty in buffer, out waits untill filled
-- FullOut: Out was filled and sent. In was not entirely read
-- Invalid[1]: Both buffers non-empty
-- Finishing: There is no more in data and in buffer is empty. Waits till
--    all outs was sent.
-- Finished: Operation finished
-- Flushing: Flush requested
-- 
-- Please note that the decompressing can finish also on flush and finish.
--
-- [1] Named for 'historical' reasons

newtype Initial = Initial BZStream
data EmptyIn = EmptyIn !BZStream !ByteString
data FullOut = FullOut !BZStream !ByteString
data Invalid = Invalid !BZStream !ByteString !ByteString
data Finishing = Finishing !BZStream !ByteString
data Flushing = Flushing !BZStream !BZipFlush !ByteString

withByteString :: ByteString -> (Ptr Word8 -> Int -> IO a) -> IO a
withByteString (PS ptr off len) f
    = withForeignPtr ptr (\ptr' -> f (ptr' `plusPtr` off) len)


{-# LINE 256 "Data/Iteratee/BZip.hsc" #-}
mkByteString :: MonadIO m => Int -> m ByteString
mkByteString s = liftIO $ create s (\_ -> return ())

{-# LINE 259 "Data/Iteratee/BZip.hsc" #-}

putOutBuffer :: Int -> BZStream -> IO ByteString
putOutBuffer size bzstr = do
    _out <- mkByteString size
    withByteString _out $ \ptr len -> withBZStream bzstr $ \bzptr -> do
        (\hsc_ptr -> pokeByteOff hsc_ptr 24) bzptr ptr
{-# LINE 265 "Data/Iteratee/BZip.hsc" #-}
        (\hsc_ptr -> pokeByteOff hsc_ptr 32) bzptr len
{-# LINE 266 "Data/Iteratee/BZip.hsc" #-}
    return _out

putInBuffer :: BZStream -> ByteString -> IO ()
putInBuffer bzstr _in
    = withByteString _in $ \ptr len -> withBZStream bzstr $ \bzptr -> do
        (\hsc_ptr -> pokeByteOff hsc_ptr 0) bzptr ptr
{-# LINE 272 "Data/Iteratee/BZip.hsc" #-}
        (\hsc_ptr -> pokeByteOff hsc_ptr 8) bzptr len
{-# LINE 273 "Data/Iteratee/BZip.hsc" #-}

pullOutBuffer :: BZStream -> ByteString -> IO ByteString
pullOutBuffer bzstr _out = withByteString _out $ \ptr _ -> do
    next_out <- withBZStream bzstr $ \bzptr -> (\hsc_ptr -> peekByteOff hsc_ptr 24) bzptr
{-# LINE 277 "Data/Iteratee/BZip.hsc" #-}
    return $! BS.take (next_out `minusPtr` ptr) _out

pullInBuffer :: BZStream -> ByteString -> IO ByteString
pullInBuffer bzstr _in = withByteString _in $ \ptr _ -> do
    next_in <- withBZStream bzstr $ \bzptr -> (\hsc_ptr -> peekByteOff hsc_ptr 0) bzptr
{-# LINE 282 "Data/Iteratee/BZip.hsc" #-}
    return $! BS.drop (next_in `minusPtr` ptr) _in

insertOut :: MonadIO m
          => Int
          -> (BZStream -> CInt -> IO CInt)
          -> Initial
          -> Enumeratee ByteString ByteString m a
insertOut size runf (Initial bzstr) iter = do
    _out <- liftIO $ putOutBuffer size bzstr

{-# LINE 294 "Data/Iteratee/BZip.hsc" #-}
    fill size runf (EmptyIn bzstr _out) iter

fill :: MonadIO m
     => Int
     -> (BZStream -> CInt -> IO CInt)
     -> EmptyIn
     -> Enumeratee ByteString ByteString m a
fill size runf (EmptyIn bzstr _out) iter
    = let fill' (Chunk _in)
              | not (BS.null _in) = do
                  liftIO $ putInBuffer bzstr _in

{-# LINE 309 "Data/Iteratee/BZip.hsc" #-}
                  doRun size runf (Invalid bzstr _in _out) iter
              | otherwise = fillI
          fill' (EOF Nothing) = do
              out <- liftIO $ pullOutBuffer bzstr _out
              iter' <- lift $ enumPure1Chunk out iter
              finish size runf (Finishing bzstr BS.empty) iter'
          fill' (EOF (Just err))
              = case fromException err of
                  Just err' -> flush size runf (Flushing bzstr err' _out) iter
                  Nothing -> throwRecoverableErr err fill'

{-# LINE 324 "Data/Iteratee/BZip.hsc" #-}
          fillI = liftI fill'

{-# LINE 326 "Data/Iteratee/BZip.hsc" #-}
      in fillI

swapOut :: MonadIO m
        => Int
        -> (BZStream -> CInt -> IO CInt)
        -> FullOut
        -> Enumeratee ByteString ByteString m a
swapOut size runf (FullOut bzstr _in) iter = do
    _out <- liftIO $ putOutBuffer size bzstr

{-# LINE 338 "Data/Iteratee/BZip.hsc" #-}
    doRun size runf (Invalid bzstr _in _out) iter

doRun :: MonadIO m
      => Int
      -> (BZStream -> CInt -> IO CInt)
      -> Invalid
      -> Enumeratee ByteString ByteString m a
doRun size runf (Invalid bzstr _in _out) iter = do

{-# LINE 350 "Data/Iteratee/BZip.hsc" #-}
    status <- liftIO $ runf bzstr 0
{-# LINE 351 "Data/Iteratee/BZip.hsc" #-}

{-# LINE 354 "Data/Iteratee/BZip.hsc" #-}
    case fromErrno status of
        Left err -> do
            _ <- joinIM $ enumErr err iter
            throwErr (toException err)
        Right False -> do -- End of stream
            remaining <- liftIO $ pullInBuffer bzstr _in
            out <- liftIO $ pullOutBuffer bzstr _out
            iter' <- lift $ enumPure1Chunk out iter
            idone iter' (Chunk remaining)
        Right True -> do -- Continue
            (avail_in, avail_out) <- liftIO $ withBZStream bzstr $ \bzptr -> do
                avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 8) bzptr
{-# LINE 366 "Data/Iteratee/BZip.hsc" #-}
                avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 32) bzptr
{-# LINE 367 "Data/Iteratee/BZip.hsc" #-}
                return (avail_in, avail_out) :: IO (CInt, CInt)
            case avail_out of
                0 -> do
                    out <- liftIO $ pullOutBuffer bzstr _out
                    iter' <- lift $ enumPure1Chunk out iter
                    case avail_in of
                        0 -> insertOut size runf (Initial bzstr) iter'
                        _ -> swapOut size runf (FullOut bzstr _in) iter'
                _ -> case avail_in of
                    0 -> fill size runf (EmptyIn bzstr _out) iter
                    _ -> do
                        _ <- joinIM $ enumErr IncorrectState iter
                        throwErr (toException IncorrectState)

flush :: MonadIO m
      => Int
      -> (BZStream -> CInt -> IO CInt)
      -> Flushing
      -> Enumeratee ByteString ByteString m a
flush size runf (Flushing bzstr _flush _out) iter = do
    status <- liftIO $ runf bzstr (fromFlush _flush)
    case fromErrno status of
        Left err -> do
            _ <- joinIM $ enumErr err iter
            throwErr (toException err)
        Right False -> do -- Finished
            out <- liftIO $ pullOutBuffer bzstr _out
            iter' <- lift $ enumPure1Chunk out iter
            idone iter' (Chunk BS.empty)
        Right True -> do
            -- TODO: avail_in unused (4th time!)
            (_avail_in, avail_out) <- liftIO $ withBZStream bzstr $ \bzptr -> do
                avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 8) bzptr
{-# LINE 400 "Data/Iteratee/BZip.hsc" #-}
                avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 32) bzptr
{-# LINE 401 "Data/Iteratee/BZip.hsc" #-}
                return (avail_in, avail_out) :: IO (CInt, CInt)
            case avail_out of
                0 -> do
                    out <- liftIO $ pullOutBuffer bzstr _out
                    iter' <- lift $ enumPure1Chunk out iter
                    out' <- liftIO $ putOutBuffer size bzstr
                    flush size runf (Flushing bzstr _flush out') iter'
                _ -> insertOut size runf (Initial bzstr) iter

finish :: MonadIO m
       => Int
       -> (BZStream -> CInt -> IO CInt)
       -> Finishing
       -> Enumeratee ByteString ByteString m a
finish size runf fin@(Finishing bzstr _in) iter = do

{-# LINE 420 "Data/Iteratee/BZip.hsc" #-}
    _out <- liftIO $ putOutBuffer size bzstr
    status <- liftIO $ runf bzstr 2
{-# LINE 422 "Data/Iteratee/BZip.hsc" #-}
    case fromErrno status of
        Left err -> do
            _ <- lift $ enumErr err iter
            throwErr (toException err)
        Right False -> do -- Finished
            remaining <- liftIO $ pullInBuffer bzstr _in
            out <- liftIO $ pullOutBuffer bzstr _out
            iter' <- lift $ enumPure1Chunk out iter
            idone iter' (Chunk remaining)
        Right True -> do
            -- TODO: avail_in is unused, is this an error or can it be removed?
            -- same thing in ZLib.hsc
            (_avail_in, avail_out) <- liftIO $ withBZStream bzstr $ \bzptr -> do
                avail_in <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 8) bzptr
{-# LINE 436 "Data/Iteratee/BZip.hsc" #-}
                avail_out <- liftIO $ (\hsc_ptr -> peekByteOff hsc_ptr 32) bzptr
{-# LINE 437 "Data/Iteratee/BZip.hsc" #-}
                return (avail_in, avail_out) :: IO (CInt, CInt)
            case avail_out of
                0 -> do
                    out <- liftIO $ pullOutBuffer bzstr _out
                    iter' <- lift $ enumPure1Chunk out iter
                    finish size runf fin iter'
                _ -> do
                    _ <-  lift $ enumErr (toException IncorrectState) iter
                    throwErr $! toException IncorrectState

foreign import ccall unsafe "BZ2_bzCompressInit"
                             compressInit :: Ptr BZStream -> CInt -> CInt
                                          -> CInt -> IO CInt
foreign import ccall unsafe "BZ2_bzDecompressInit"
                             decompressInit :: Ptr BZStream -> CInt -> CInt
                                            -> IO CInt
foreign import ccall unsafe "BZ2_bzCompress"
                             compress :: Ptr BZStream -> CInt -> IO CInt
foreign import ccall unsafe "BZ2_bzDecompress"
                             decompress :: Ptr BZStream -> IO CInt
foreign import ccall unsafe "&BZ2_bzCompressEnd"
                             compressEnd :: FunPtr (Ptr BZStream -> IO ())
foreign import ccall unsafe "&BZ2_bzDecompressEnd"
                             decompressEnd :: FunPtr (Ptr BZStream -> IO ())


{-# LINE 476 "Data/Iteratee/BZip.hsc" #-}
compress' :: BZStream -> CInt -> IO CInt
compress' bz f = withBZStream bz $ \p -> compress p f

decompress' :: BZStream -> CInt -> IO CInt
decompress' bz _ = withBZStream bz decompress

verboseLevel :: CInt
verboseLevel = 0

{-# LINE 485 "Data/Iteratee/BZip.hsc" #-}

mkCompress :: CompressParams -> IO (Either BZipParamsException Initial)
mkCompress (CompressParams blk wf _)
    = case fromBlockSize blk of
        Left err -> return $! Left $! err
        Right blk' -> case fromWorkFactor wf of
            Left err -> return $! Left $! err
            Right wf' -> do
                bzstr <- mallocForeignPtrBytes (80)
{-# LINE 494 "Data/Iteratee/BZip.hsc" #-}
                withForeignPtr bzstr $ \bzptr -> do
                    memset (castPtr bzptr) 0 (80)
{-# LINE 496 "Data/Iteratee/BZip.hsc" #-}
                    compressInit bzptr blk' verboseLevel wf' `finally`
                        addForeignPtrFinalizer compressEnd bzstr
                return $! Right $! Initial $ BZStream bzstr

mkDecompress :: DecompressParams -> IO (Either BZipParamsException Initial)
mkDecompress (DecompressParams small _) = do
    bzstr <- mallocForeignPtrBytes (80)
{-# LINE 503 "Data/Iteratee/BZip.hsc" #-}
    withForeignPtr bzstr $ \bzptr -> do
        memset (castPtr bzptr) 0 (80)
{-# LINE 505 "Data/Iteratee/BZip.hsc" #-}
        decompressInit bzptr verboseLevel (if small then 0 else 1) `finally`
            addForeignPtrFinalizer decompressEnd bzstr
    return $! Right $! Initial $ BZStream bzstr

-- User-related code

-- | Compress the input and send to inner iteratee.
enumCompress :: MonadIO m
             => CompressParams -- ^ Parameters of compression
             -> Enumeratee ByteString ByteString m a
enumCompress cp@(CompressParams _ _ size) iter = do
    cmp <- liftIO $ mkCompress cp
    case cmp of
        Left err -> do
            _ <- lift $ enumErr err iter
            throwErr (toException err)
        Right init' -> insertOut size compress' init' iter

-- | Decompress the input and send to inner iteratee. If there is end of
-- zlib stream it is left unprocessed.
enumDecompress :: MonadIO m
               => DecompressParams
               -> Enumeratee ByteString ByteString m a
enumDecompress dp@(DecompressParams _ size) iter = do
    dcmp <- liftIO $ mkDecompress dp
    case dcmp of
        Left err -> do
            _ <- lift $ enumErr err iter
            throwErr (toException err)
        Right init' -> insertOut size decompress' init' iter