{-# LINE 1 "Z/Compression/Zlib.hsc" #-}
{-|
Module      : Z.Compression.Zlib
Description : The zlib
Copyright   : (c) Dong Han, 2017-2018
License     : BSD
Maintainer  : winterland1989@gmail.com
Stability   : experimental
Portability : non-portable
This module provides <https://zlib.net zlib> bindings.
-}

module Z.Compression.Zlib(
  -- * Compression
    CompressConfig(..)
  , defaultCompressConfig
  , compress
  , compressSink
  , Strategy (Z_FILTERED, Z_HUFFMAN_ONLY, Z_RLE, Z_FIXED, Z_DEFAULT_STRATEGY)
  , CompressLevel(Z_BEST_SPEED, Z_BEST_COMPRESSION, Z_DEFAULT_COMPRESSION)
  , WindowBits
  , defaultWindowBits
  , MemLevel
  , defaultMemLevel
  -- * Decompression
  , DecompressConfig(..)
  , defaultDecompressConfig
  , decompress
  , decompressSource
  ) where

import           Control.Monad
import           Data.IORef
import           Data.Typeable
import           Data.Word
import           Foreign            hiding (void)
import           Foreign.C
import           GHC.Generics
import           System.IO.Unsafe   (unsafePerformIO)
import           Z.Data.Array       as A
import           Z.Data.CBytes      as CBytes
import           Z.Data.Vector.Base as V
import           Z.Data.Text.ShowT  (ShowT)
import           Z.Foreign
import           Z.IO.Buffered
import           Z.IO.Exception



newtype Strategy = Strategy CInt deriving (Eq, Ord, Show, Generic)
                                    deriving anyclass ShowT

pattern Z_FILTERED           :: Strategy
pattern Z_HUFFMAN_ONLY       :: Strategy
pattern Z_RLE                :: Strategy
pattern Z_FIXED              :: Strategy
pattern Z_DEFAULT_STRATEGY   :: Strategy
pattern Z_FILTERED           = Strategy (1)
{-# LINE 58 "Z/Compression/Zlib.hsc" #-}
pattern Z_HUFFMAN_ONLY       = Strategy (2)
{-# LINE 59 "Z/Compression/Zlib.hsc" #-}
pattern Z_RLE                = Strategy (3)
{-# LINE 60 "Z/Compression/Zlib.hsc" #-}
pattern Z_FIXED              = Strategy (4)
{-# LINE 61 "Z/Compression/Zlib.hsc" #-}
pattern Z_DEFAULT_STRATEGY   = Strategy (0)
{-# LINE 62 "Z/Compression/Zlib.hsc" #-}


newtype CompressLevel = CompressLevel CInt deriving (Eq, Ord, Show, Generic)
                                            deriving anyclass ShowT

-- pattern Z_NO_COMPRESSION       =  CompressLevel (#const Z_NO_COMPRESSION     )
pattern Z_BEST_SPEED          :: CompressLevel
pattern Z_BEST_COMPRESSION    :: CompressLevel
pattern Z_DEFAULT_COMPRESSION :: CompressLevel
pattern Z_BEST_SPEED           =  CompressLevel (1)
{-# LINE 72 "Z/Compression/Zlib.hsc" #-}
pattern Z_BEST_COMPRESSION     =  CompressLevel (9)
{-# LINE 73 "Z/Compression/Zlib.hsc" #-}
pattern Z_DEFAULT_COMPRESSION  =  CompressLevel (-1)
{-# LINE 74 "Z/Compression/Zlib.hsc" #-}

{- | The 'WindowBits' is the base two logarithm of the maximum window size (the size of the history buffer).
It should be in the range 8..15 for this version of the library. The 'defaultWindowBits' value is 15. Decompressing windowBits must be greater than or equal to the compressing windowBits. If a compressed stream with a larger window size is given as input, decompress will throw 'ZDataError'
windowBits can also be –8..–15 for raw inflate. In this case, -windowBits determines the window size. inflate() will then process raw deflate data, not looking for a zlib or gzip header, not generating a check value, and not looking for any check values for comparison at the end of the stream.
windowBits can also be greater than 15 for optional gzip decoding. Add 32 to windowBits to enable zlib and gzip decoding with automatic header detection, or add 16 to decode only the gzip format.
-}
newtype WindowBits = WindowBits CInt
    deriving (Eq, Ord, Read, Show, Generic)
        deriving newtype Num
        deriving anyclass ShowT

defaultWindowBits :: WindowBits
defaultWindowBits = WindowBits 15

-- | The 'MemLevel' specifies how much memory should be allocated for the internal compression state. 1 uses minimum memory but is slow and reduces compression ratio; 9 uses maximum memory for optimal speed. The default value is 8.
newtype MemLevel = MemLevel CInt
    deriving (Eq, Ord, Read, Show, Generic)
        deriving newtype Num
        deriving anyclass ShowT

defaultMemLevel :: MemLevel
defaultMemLevel = MemLevel 9

data CompressConfig = CompressConfig
    { compressLevel :: CompressLevel
    , compressWindowBits :: WindowBits
    , compressMemoryLevel :: MemLevel
    , compressDictionary :: V.Bytes
    , compressStrategy :: Strategy
    }

defaultCompressConfig :: CompressConfig
defaultCompressConfig =
    CompressConfig Z_DEFAULT_COMPRESSION  defaultWindowBits
        defaultMemLevel V.empty Z_DEFAULT_STRATEGY

-- | Compress all the data written to a output.
--
compressSink :: HasCallStack
           => CompressConfig
           -> Sink V.Bytes
           -> IO (Sink V.Bytes)
compressSink (CompressConfig level windowBits memLevel dict strategy) (write, flush) = do
    zs <- newForeignPtr free_z_stream_deflate =<< create_z_stream
    buf <- A.newPinnedPrimArray bufSiz
    set_avail_out zs buf bufSiz
    bufRef <- newIORef buf

    withForeignPtr zs $ \ ps -> do
        throwZlibIfMinus_ $ deflate_init2 ps level windowBits memLevel strategy
        unless (V.null dict) $
            throwZlibIfMinus_ . withPrimVectorUnsafe dict $ \ pdict off len ->
            deflate_set_dictionary ps pdict off (fromIntegral $ len)

    return (zwrite zs bufRef, zflush zs bufRef)

  where
    bufSiz = V.defaultChunkSize

    zwrite zs bufRef input = do
        set_avail_in zs input (V.length input)
        zloop zs bufRef

    zloop zs bufRef = do
        oavail :: CUInt <- withForeignPtr zs $ \ ps -> do
            throwZlibIfMinus_ (deflate ps (0))
{-# LINE 140 "Z/Compression/Zlib.hsc" #-}
            ((\hsc_ptr -> peekByteOff hsc_ptr 32)) ps
{-# LINE 141 "Z/Compression/Zlib.hsc" #-}

        when (oavail == 0) $ do
            oarr <- A.unsafeFreezeArr =<< readIORef bufRef
            buf' <- A.newPinnedPrimArray bufSiz
            set_avail_out zs buf' bufSiz
            writeIORef bufRef buf'
            write (V.PrimVector oarr 0 bufSiz)
            zloop zs bufRef

    zflush zs bufRef = do
        r :: CInt <- withForeignPtr zs $ \ ps -> do
            r <- throwZlibIfMinus (deflate ps (4))
{-# LINE 153 "Z/Compression/Zlib.hsc" #-}
            oavail :: CUInt <- ((\hsc_ptr -> peekByteOff hsc_ptr 32)) ps
{-# LINE 154 "Z/Compression/Zlib.hsc" #-}
            when (oavail /= fromIntegral bufSiz) $ do
                oarr <- A.unsafeFreezeArr =<< readIORef bufRef
                write (V.PrimVector oarr 0 (bufSiz - fromIntegral oavail))
                flush
            return r

        when (r /= (1)) $ do
{-# LINE 161 "Z/Compression/Zlib.hsc" #-}
            buf' <- A.newPinnedPrimArray bufSiz
            set_avail_out zs buf' bufSiz
            writeIORef bufRef buf'
            zflush zs bufRef

-- | Compress some bytes.
compress :: HasCallStack => CompressConfig -> V.Bytes -> V.Bytes
compress conf input = unsafePerformIO $ do
    ref <- newIORef []
    (write, flush) <- compressSink conf (\ x -> modifyIORef' ref (x:), return ())
    write input
    flush
    V.concat . reverse <$> readIORef ref


{-
compressBuilderStream :: HasCallStack
                      => CompressConfig
                      -> (B.Builder a -> IO ())
                      -> IO (B.Builder a -> IO ())


-}

data DecompressConfig = DecompressConfig
    { decompressWindowBits :: WindowBits
    , decompressDictionary :: V.Bytes
    }

defaultDecompressConfig :: DecompressConfig
defaultDecompressConfig = DecompressConfig defaultWindowBits V.empty

-- | Decompress bytes from source.
decompressSource :: DecompressConfig
                 -> Source V.Bytes
                 -> IO (Source V.Bytes)
decompressSource (DecompressConfig windowBits dict) source = do
    zs <- newForeignPtr free_z_stream_inflate =<< create_z_stream
    buf <- A.newPinnedPrimArray bufSiz
    set_avail_out zs buf bufSiz
    bufRef <- newIORef buf

    withForeignPtr zs $ \ ps -> do
        throwZlibIfMinus_ $ inflate_init2 ps windowBits

    return (zread zs bufRef)
  where
    bufSiz = V.defaultChunkSize

    zread zs bufRef = do
        bufLen <- A.sizeofMutableArr =<< readIORef bufRef
        if bufLen == 0
        then return Nothing
        else do
            oavail :: CUInt <- withForeignPtr zs ((\hsc_ptr -> peekByteOff hsc_ptr 32))
{-# LINE 216 "Z/Compression/Zlib.hsc" #-}
            if (oavail == 0)
            then do
                oarr <- A.unsafeFreezeArr =<< readIORef bufRef
                buf' <- A.newPinnedPrimArray bufSiz
                set_avail_out zs buf' bufSiz
                writeIORef bufRef buf'
                return (Just (V.PrimVector oarr 0 bufSiz))
            else zloop zs bufRef

    zloop zs bufRef  = do
        iavail :: CUInt <- withForeignPtr zs ((\hsc_ptr -> peekByteOff hsc_ptr 8))
{-# LINE 227 "Z/Compression/Zlib.hsc" #-}
        if iavail == 0
        then do
            input <- source
            case input of
                Just input' -> do
                    set_avail_in zs input' (V.length input')
                    withForeignPtr zs $ \ ps -> do
                        r <- throwZlibIfMinus (inflate ps (0))
{-# LINE 235 "Z/Compression/Zlib.hsc" #-}
                        when (r == (2) && not (V.null dict)) $ do
{-# LINE 236 "Z/Compression/Zlib.hsc" #-}
                            throwZlibIfMinus_ . withPrimVectorUnsafe dict $ \ pdict off len ->
                                inflate_set_dictionary ps pdict off (fromIntegral len)
                    zread zs bufRef
                _ -> zfinish zs bufRef []
        else do
            withForeignPtr zs $ \ ps ->
                throwZlibIfMinus_ (inflate ps (0))
{-# LINE 243 "Z/Compression/Zlib.hsc" #-}
            zloop zs bufRef

    zfinish zs bufRef acc = do
        r <- withForeignPtr zs $ \ ps -> do
            throwZlibIfMinus (inflate ps (4))
{-# LINE 248 "Z/Compression/Zlib.hsc" #-}

        oavail :: CUInt <- withForeignPtr zs ((\hsc_ptr -> peekByteOff hsc_ptr 32))
{-# LINE 250 "Z/Compression/Zlib.hsc" #-}
        oarr <- A.unsafeFreezeArr =<< readIORef bufRef
        let !v = V.PrimVector oarr 0 (bufSiz - fromIntegral oavail)

        if (r == (1))
{-# LINE 254 "Z/Compression/Zlib.hsc" #-}
        then do
            writeIORef bufRef =<< A.newArr 0
            let !v' = V.concat (reverse (v:acc))
            return (Just v')
        else do
            buf' <- A.newPinnedPrimArray bufSiz
            set_avail_out zs buf' bufSiz
            writeIORef bufRef buf'
            zfinish zs bufRef (v:acc)


-- | Decompress some bytes.
decompress :: HasCallStack => DecompressConfig -> V.Bytes -> V.Bytes
decompress conf input = V.concat . unsafePerformIO $ do
     collectSource =<< decompressSource conf =<< sourceFromList [input]

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

newtype ZReturn = ZReturn CInt deriving (Eq, Ord, Show, Typeable)

toZErrorMsg :: CInt -> CBytes
toZErrorMsg (0) =  "Z_OK"
{-# LINE 276 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (1) =  "Z_STREAM_END"
{-# LINE 277 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (2) =  "Z_NEED_DICT"
{-# LINE 278 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-1) =  "Z_ERRNO"
{-# LINE 279 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-2) =  "Z_STREAM_ERROR"
{-# LINE 280 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-3) =  "Z_DATA_ERROR"
{-# LINE 281 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-4) =  "Z_MEM_ERROR"
{-# LINE 282 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-5) =  "Z_BUF_ERROR"
{-# LINE 283 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg (-6) =  "Z_VERSION_ERROR"
{-# LINE 284 "Z/Compression/Zlib.hsc" #-}
toZErrorMsg _                        =  "Z_UNEXPECTED"

data ZlibException = ZlibException CBytes CallStack deriving (Show, Typeable)
instance Exception ZlibException

throwZlibIfMinus :: HasCallStack => IO CInt -> IO CInt
throwZlibIfMinus f = do
    r <- f
    if r < 0 && r /= (-5)
{-# LINE 293 "Z/Compression/Zlib.hsc" #-}
    then throwIO (ZlibException (toZErrorMsg r) callStack)
    else return r

throwZlibIfMinus_ :: HasCallStack => IO CInt -> IO ()
throwZlibIfMinus_ = void . throwZlibIfMinus

data ZStream

foreign import ccall unsafe
    create_z_stream :: IO (Ptr ZStream)

foreign import ccall unsafe "hs_zlib.c &free_z_stream_inflate"
    free_z_stream_inflate :: FunPtr (Ptr ZStream -> IO ())

foreign import ccall unsafe "hs_zlib.c &free_z_stream_deflate"
    free_z_stream_deflate :: FunPtr (Ptr ZStream -> IO ())

foreign import ccall unsafe
    deflate_init2 :: Ptr ZStream -> CompressLevel -> WindowBits -> MemLevel -> Strategy -> IO CInt

foreign import ccall unsafe
    deflate_set_dictionary :: Ptr ZStream -> BA# Word8 -> Int -> Int -> IO CInt

foreign import ccall unsafe
    deflate :: Ptr ZStream -> CInt -> IO CInt

foreign import ccall unsafe
    inflate_init2 :: Ptr ZStream -> WindowBits -> IO CInt

foreign import ccall unsafe
    inflate_set_dictionary :: Ptr ZStream -> BA# Word8 -> Int -> Int -> IO CInt

foreign import ccall unsafe
    inflate :: Ptr ZStream -> CInt -> IO CInt

set_avail_in :: ForeignPtr ZStream -> V.Bytes -> Int -> IO ()
set_avail_in zs buf buflen = do
    withPrimVectorSafe buf $ \ pbuf _ ->
        withForeignPtr zs $ \ ps -> do
            ((\hsc_ptr -> pokeByteOff hsc_ptr 0)) ps pbuf
{-# LINE 333 "Z/Compression/Zlib.hsc" #-}
            ((\hsc_ptr -> pokeByteOff hsc_ptr 8)) ps (fromIntegral buflen :: CUInt)
{-# LINE 334 "Z/Compression/Zlib.hsc" #-}

set_avail_out :: ForeignPtr ZStream -> MutablePrimArray RealWorld Word8 -> Int -> IO ()
set_avail_out zs buf bufSiz = do
    withMutablePrimArrayContents buf $ \ pbuf ->
        withForeignPtr zs $ \ ps -> do
            ((\hsc_ptr -> pokeByteOff hsc_ptr 24)) ps pbuf
{-# LINE 340 "Z/Compression/Zlib.hsc" #-}
            ((\hsc_ptr -> pokeByteOff hsc_ptr 32)) ps (fromIntegral bufSiz :: CUInt)
{-# LINE 341 "Z/Compression/Zlib.hsc" #-}