    -- |
    -- Module      : Data.Git.Storage.FileReader
    -- License     : BSD-style
    -- Maintainer  : Vincent Hanquez <vincent@snarc.org>
    -- Stability   : experimental
    -- Portability : unix
    --
    {-# LANGUAGE DeriveDataTypeable #-}
    module Data.Git.Storage.FileReader
            ( FileReader
            , fileReaderNew
            , fileReaderClose
            , withFileReader
            , withFileReaderDecompress
            , fileReaderGetPos
            , fileReaderGet
            , fileReaderGetLBS
            , fileReaderGetBS
            , fileReaderGetVLF
            , fileReaderSeek
            , fileReaderParse

            , fileReaderInflateToSize
            ) where

    import Control.Applicative ((<$>))
    import Control.Exception (bracket, throwIO)
    import Control.Monad

    import Data.Attoparsec (parseWith, Parser, IResult(..))
    import qualified Data.Attoparsec as A
    import Data.Bits
    import Data.ByteString (ByteString)
    import Data.ByteString.Unsafe
    import qualified Data.ByteString as B
    import qualified Data.ByteString.Lazy as L
    import Data.ByteString.Lazy.Internal (defaultChunkSize)
    import Data.IORef

    import Data.Data
    import Data.Word

    import Codec.Zlib
    import Codec.Zlib.Lowlevel
    import Foreign.ForeignPtr
    import qualified Control.Exception as E

    import System.IO (hClose, hSeek, SeekMode(..))
    import Filesystem
    import Filesystem.Path
    import Prelude hiding (FilePath)

    data FileReader = FileReader
            { fbHandle     :: Handle
            , fbUseInflate :: Bool
            , fbInflate    :: Inflate
            , fbRemaining  :: IORef ByteString
            , fbPos        :: IORef Word64
            }

    data InflateException = InflateException Word64 Word64 String
        deriving (Show,Eq,Typeable)

    instance E.Exception InflateException

    fileReaderNew :: Bool -> Handle -> IO FileReader
    fileReaderNew decompress handle = do
            ref <- newIORef B.empty
            pos <- newIORef 0
            inflate <- initInflate defaultWindowBits
            return $ FileReader handle decompress inflate ref pos

    fileReaderClose :: FileReader -> IO ()
    fileReaderClose = hClose . fbHandle

    withFileReader :: FilePath -> (FileReader -> IO a) -> IO a
    withFileReader path f =
            bracket (openFile path ReadMode) (hClose) $ \handle ->
                    bracket (fileReaderNew False handle) (\_ -> return ()) f

    withFileReaderDecompress :: FilePath -> (FileReader -> IO a) -> IO a
    withFileReaderDecompress path f =
            bracket (openFile path ReadMode) (hClose) $ \handle ->
                    bracket (fileReaderNew True handle) (\_ -> return ()) f

    fileReaderGetNext :: FileReader -> IO ByteString
    fileReaderGetNext fb = do
            bs <- if fbUseInflate fb then inflateTillPop else B.hGet (fbHandle fb) 8192
            modifyIORef (fbPos fb) (\pos -> pos + (fromIntegral $ B.length bs))
            return bs
            where inflateTillPop = do
                    b <- B.hGet (fbHandle fb) 4096
                    if B.null b
                            then finishInflate (fbInflate fb)
                            else (>>= maybe inflateTillPop return) =<< feedInflate (fbInflate fb) b

    fileReaderGetPos :: FileReader -> IO Word64
    fileReaderGetPos fr = do
            storeLeft <- B.length <$> readIORef (fbRemaining fr)
            pos       <- readIORef (fbPos fr)
            return (pos - fromIntegral storeLeft)

    fileReaderFill :: FileReader -> IO ()
    fileReaderFill fb = fileReaderGetNext fb >>= writeIORef (fbRemaining fb)

    fileReaderGet :: Int -> FileReader -> IO [ByteString]
    fileReaderGet size fb@(FileReader { fbRemaining = ref }) = loop size
            where loop left = do
                    b <- readIORef ref
                    if B.length b >= left
                            then do
                                    let (b1, b2) = B.splitAt left b
                                    writeIORef ref b2
                                    return [b1]
                            else do
                                    let nleft = left - B.length b
                                    fileReaderFill fb
                                    liftM (b :) (loop nleft)

    fileReaderGetLBS :: Int -> FileReader -> IO L.ByteString
    fileReaderGetLBS size fb = L.fromChunks <$> fileReaderGet size fb

    fileReaderGetBS :: Int -> FileReader -> IO ByteString
    fileReaderGetBS size fb = B.concat <$> fileReaderGet size fb

    -- | seek in a handle, and reset the remaining buffer to empty.
    fileReaderSeek :: FileReader -> Word64 -> IO ()
    fileReaderSeek (FileReader { fbHandle = handle, fbRemaining = ref, fbPos = pos }) absPos = do
            writeIORef ref B.empty >> writeIORef pos absPos >> hSeek handle AbsoluteSeek (fromIntegral absPos)

    -- | parse from a filebuffer
    fileReaderParse :: FileReader -> Parser a -> IO a
    fileReaderParse fr@(FileReader { fbRemaining = ref }) parseF = do
            initBS <- readIORef ref
            result <- parseWith (fileReaderGetNext fr) parseF initBS
            case result of
                    Done remaining a -> writeIORef ref remaining >> return a
                    Partial _        -> error "parsing failed: partial with a handle, reached EOF ?"
                    Fail _ ctxs err  -> error ("parsing failed: " ++ show ctxs ++ " : " ++ show err)

    -- | get a Variable Length Field. get byte as long as MSB is set, and one byte after
    fileReaderGetVLF :: FileReader -> IO [Word8]
    fileReaderGetVLF fr = fileReaderParse fr $ do
            bs <- A.takeWhile (\w -> w `testBit` 7)
            l  <- A.anyWord8
            return $ (map (\w -> w `clearBit` 7) $ B.unpack bs) ++ [l]

    fileReaderInflateToSize :: FileReader -> Word64 -> IO L.ByteString
    fileReaderInflateToSize fb@(FileReader { fbRemaining = ref }) outputSize = do
            --pos <- fileReaderGetPos fb
            --putStrLn ("inflate to size " ++ show outputSize ++ " " ++ show pos)
            inflate <- inflateNew
            l       <- loop inflate outputSize
            --posend <- fileReaderGetPos fb
            --putStrLn ("inflated input " ++ show posend)
            return $ L.fromChunks l
            where loop inflate left = do
                    rbs <- readIORef ref
                    let maxToInflate = min left (16 * 1024)
                    let lastBlock = if left == maxToInflate then True else False
                    (dbs,remaining) <- inflateToSize inflate (fromIntegral maxToInflate) lastBlock rbs (fileReaderGetNext fb)
                                       `E.catch` augmentAndRaise left
                    writeIORef ref remaining
                    let nleft = left - fromIntegral (B.length dbs)
                    if nleft > 0
                            then liftM (dbs:) (loop inflate nleft)
                            else return [dbs]
                  augmentAndRaise :: Word64 -> E.SomeException -> IO a
                  augmentAndRaise left exn = throwIO $ InflateException outputSize left (show exn)

    -- lowlevel helpers to inflate only to a specific size.

    inflateNew = do
            zstr <- zstreamNew
            inflateInit2 zstr defaultWindowBits
            newForeignPtr c_free_z_stream_inflate zstr

    inflateToSize inflate sz isLastBlock ibs nextBs = withForeignPtr inflate $ \zstr -> do
            let boundSz = min defaultChunkSize sz
            -- create an output buffer
            fbuff <- mallocForeignPtrBytes boundSz
            withForeignPtr fbuff $ \buff -> do
                    c_set_avail_out zstr buff (fromIntegral boundSz)
                    rbs <- loop zstr ibs
                    bs  <- B.packCStringLen (buff, boundSz)
                    return (bs, rbs)
            where
                    loop zstr nbs = do
                            (ai, streamEnd) <- inflateOneInput zstr nbs
                            ao <- c_get_avail_out zstr
                            if (isLastBlock && streamEnd) || (not isLastBlock && ao == 0)
                                    then return $ bsTakeLast ai nbs
                                    else do
                                            --when (ai /= 0) $ error ("input not consumed completly: ai" ++ show ai)
                                            (if ai == 0
                                                    then nextBs
                                                    else return (bsTakeLast ai nbs)) >>= loop zstr

                    inflateOneInput zstr bs = unsafeUseAsCStringLen bs $ \(istr, ilen) -> do
                            c_set_avail_in zstr istr $ fromIntegral ilen
                            r <- c_call_inflate_noflush zstr
                            when (r < 0 && r /= (-5)) $ do
                                    throwIO $ ZlibException $ fromIntegral r
                            ai <- c_get_avail_in zstr
                            return (ai, r == 1)
                    bsTakeLast len bs = B.drop (B.length bs - fromIntegral len) bs