    module Data.Binary.Defer.Monad(
        DeferPut, putDefer, runDeferPut,
        putInt, putByte, putChr, putByteString,
        DeferGet, getDefer, runDeferGet,
        getInt, getByte, getChr, getByteString,
        getDeferGet, getDeferPut
        ) where

    import System.IO
    import System.IO.Unsafe
    import Data.Binary.Raw
    import Control.Monad.Reader
    import Data.IORef
    import qualified Data.ByteString as BS

    import Data.Typeable
    import qualified Data.TypeMap as TypeMap
    import Foreign

    ---------------------------------------------------------------------
    -- Defer Put

    -- Storing the position explicitly gives a ~5% speed up
    -- and removes hGetPos as being a bottleneck
    -- possibly still not worth it though

    type DeferPut a = ReaderT (Buffer, IORef [DeferPending], IORef [DeferPatchup]) IO a
    data DeferPending = DeferPending !Int (DeferPut ())
    data DeferPatchup = DeferPatchup !Int !Int -- a b = at position a, write out b

    putValue :: Storable a => a -> DeferPut ()
    putValue x = do
        (buf,_,_) <- ask
        lift $ bufferAdd buf x

    putInt :: Int -> DeferPut ()
    putInt x = putValue (fromIntegral x :: Int32)

    putByte :: Word8 -> DeferPut ()
    putByte x = putValue x

    putChr :: Char -> DeferPut ()
    putChr  x = putByte $ fromIntegral $ fromEnum x

    putByteString :: BS.ByteString -> DeferPut ()
    putByteString x = do
        (buf,_,_) <- ask
        putInt $ BS.length x
        lift $ bufferAddByteString buf x

    putDefer :: DeferPut () -> DeferPut ()
    putDefer x = do
        (buf,ref,_) <- ask
        lift $ do
            p <- bufferPos buf
            bufferAdd buf (0 :: Int32) -- to backpatch
            modifyIORef ref (DeferPending p x :)

    runDeferPut :: Handle -> DeferPut () -> IO ()
    runDeferPut h m = do
        buf <- bufferNew h
        ref <- newIORef []
        back <- newIORef []
        runReaderT (m >> runDeferPendings) (buf,ref,back)
        bufferFlush buf
        patch <- readIORef back
        mapM_ (\(DeferPatchup a b) -> do hSetPos h (toInteger a); hPutInt h b) patch

    runDeferPendings :: DeferPut ()
    runDeferPendings = do
        (_,ref,back) <- ask
        todo <- lift $ readIORef ref
        lift $ writeIORef ref []
        mapM_ runDeferPending todo

    runDeferPending :: DeferPending -> DeferPut ()
    runDeferPending (DeferPending pos act) = do
        (buf,_,back) <- ask
        lift $ do
            p <- bufferPos buf
            b <- bufferPatch buf pos (fromIntegral p :: Int32)
            unless b $ modifyIORef back (DeferPatchup pos p :)
        act
        runDeferPendings

    ---------------------------------------------------------------------
    -- Buffer for writing

    bufferSize = 10000 :: Int

    -- (number in file, number in buffer)
    data Buffer = Buffer !Handle !(IORef Int) !(Ptr ()) !(IORef Int)

    bufferNew :: Handle -> IO Buffer
    bufferNew h = do
        i <- hGetPos h
        file <- newIORef $ fromInteger i
        buf <- newIORef 0
        ptr <- mallocBytes bufferSize
        return $ Buffer h file ptr buf

    bufferAdd :: Storable a => Buffer -> a -> IO ()
    bufferAdd (Buffer h file ptr buf) x = do
        let sz = sizeOf x
        buf2 <- readIORef buf
        if sz + buf2 >= bufferSize then do
            hPutBuf h ptr buf2
            pokeByteOff ptr 0 x
            modifyIORef file (+buf2)
            writeIORef buf sz
         else do
            pokeByteOff ptr buf2 x
            writeIORef buf (buf2+sz)

    bufferAddByteString :: Buffer -> BS.ByteString -> IO ()
    bufferAddByteString (Buffer h file ptr buf) x = do
        let sz = BS.length x
        buf2 <- readIORef buf
        when (buf2 /= 0) $ do
            hPutBuf h ptr buf2
            writeIORef buf 0
        modifyIORef file (+ (buf2+sz))
        BS.hPut h x

    bufferFlush :: Buffer -> IO ()
    bufferFlush (Buffer h file ptr buf) = do
        buf2 <- readIORef buf
        hPutBuf h ptr buf2
        modifyIORef file (+buf2)
        writeIORef buf 0

    bufferPos :: Buffer -> IO Int
    bufferPos (Buffer h file ptr buf) = do
        i <- readIORef file
        j <- readIORef buf
        return $ i + j

    -- Patch at position p, with value v
    -- Return True if you succeeded
    bufferPatch :: Buffer -> Int -> Int32 -> IO Bool
    bufferPatch (Buffer h file ptr buf) p v = do
        i <- readIORef file
        if p < i then return False else do
            pokeByteOff ptr (p-i) v
            return True

    ---------------------------------------------------------------------
    -- Defer Get

    type DeferGet a = ReaderT (Handle, IORef TypeMap.TypeMap) IO a

    getInt :: DeferGet Int
    getInt  = do h <- asks fst; lift $ hGetInt  h

    getByte :: DeferGet Word8
    getByte = do h <- asks fst; lift $ liftM fromIntegral $ hGetByte h

    getChr :: DeferGet Char
    getChr  = do h <- asks fst; lift $ hGetChar h

    getByteString :: DeferGet BS.ByteString
    getByteString = do
        h <- asks fst
        len <- lift $ hGetInt h
        lift $ BS.hGet h len

    getDefer :: DeferGet a -> DeferGet a
    getDefer x = do
        h <- asks fst
        i <- lift $ hGetInt h
        s <- ask
        lift $ unsafeInterleaveIO $ do
            hSetPos h (toInteger i)
            runReaderT x s

    runDeferGet :: Handle -> DeferGet a -> IO a
    runDeferGet h m = do
        ref <- newIORef TypeMap.empty
        runReaderT m (h,ref)

    getDeferGet :: Typeable a => DeferGet a
    getDeferGet = do
        ref <- asks snd
        mp <- lift $ readIORef ref
        return $ TypeMap.find mp

    getDeferPut :: Typeable a => a -> DeferGet ()
    getDeferPut x = do
        ref <- asks snd
        lift $ modifyIORef ref $ TypeMap.insert x