content: -- | This module provides helpers for emitting and reading binary files with
  -- a trailing "header".
  module BTree.BinaryFile
      ( writeWithHeader
      , readWithHeader
      ) where
  
  import Control.Monad (when)
  import Control.Error
  import Control.Monad.Trans.Class
  import Control.Monad.Catch
  import Control.Applicative
  import Data.Word
  import System.IO
  import Prelude
  
  import qualified Data.ByteString.Lazy as LBS
  import qualified Data.Binary as B
  import qualified Data.Binary.Get as B
  import qualified Data.Binary.Put as B
  import Pipes
  
  -- | An internal data structure placed at the very end of the file which
  -- describes the header and provides a magic number for sanity checking.
  data Epilogue = Epilogue { magic :: Word64
                           , headerLen :: Word64
                           }
                deriving (Show)
  
  epiLength :: Integer
  epiLength = 16
  
  magicNumber :: Word64
  magicNumber = 0xdeadbeef
  
  instance B.Binary Epilogue where
      get = Epilogue <$> B.getWord64le <*> B.getWord64le
      put (Epilogue m l) = B.putWord64le m >> B.putWord64le l
  
  -- | Write the produced 'LBS.ByteString's to the file followed by the
  -- returned header
  writeWithHeader :: (MonadMask m, MonadIO m, B.Binary hdr)
                  => FilePath
                  -> Producer LBS.ByteString m (hdr, r)
                  -> m r
  writeWithHeader fname prod =
      bracket (liftIO $ openFile fname WriteMode) (liftIO . hClose)
      $ \hdl -> hWriteWithHeader hdl prod
  
  -- | Write the produced 'LBS.ByteString's to the file followed by the
  -- returned header
  hWriteWithHeader :: (MonadIO m, B.Binary hdr)
                  => Handle
                  -> Producer LBS.ByteString m (hdr, r)
                  -> m r
  hWriteWithHeader h prod = do
      (hdr, r) <- runEffect $ for prod (liftIO . LBS.hPut h)
      let encoded = B.encode hdr
      liftIO $ LBS.hPut h encoded
      let epi = Epilogue { magic = magicNumber
                         , headerLen = fromIntegral $ LBS.length encoded }
      liftIO $ LBS.hPut h (B.encode epi)
      return r
  {-# INLINE writeWithHeader #-}
  
  annotate :: Monad m => String -> ExceptT String m a -> ExceptT String m a
  annotate ann = fmapLT ((ann++": ")++)
  
  runGetT :: Monad m => B.Get a -> LBS.ByteString -> ExceptT String m a
  runGetT _get bs =
      case B.runGetOrFail _get bs of
        Left (_, _, e)  -> throwE e
        Right (_, _, a) -> return a
  
  -- | Read and verify the header from the file, then pass it along with the
  -- file's handle to an action. The file handle sits at the beginning of the
  -- written content when passed to the action.
  readWithHeader :: (MonadMask m, MonadIO m, B.Binary hdr)
                 => FilePath
                 -> (hdr -> Handle -> m a)
                 -> ExceptT String m a
  readWithHeader fname action = do
      r <- lift $ bracket (liftIO $ openFile fname ReadMode) (liftIO . hClose) $ \h -> runExceptT $ do
          -- read epilogue
          liftIO $ hSeek h SeekFromEnd (-epiLength)
          epiBytes <- liftIO (LBS.hGet h $ fromIntegral epiLength)
          epi <- annotate "Error reading epilogue" (runGetT B.get epiBytes)
          when (magic epi /= magicNumber) $
              throwE "BinaryFile.readWithHeader: Bad magic number"
          -- read header
          let offset = fromIntegral epiLength + fromIntegral (headerLen epi)
          liftIO $ hSeek h SeekFromEnd (negate offset)
          hdrBytes <- liftIO (LBS.hGet h $ fromIntegral $ headerLen epi)
          hdr <- annotate "Error reading header" (runGetT B.get hdrBytes)
          -- pass control to action
          liftIO $ hSeek h AbsoluteSeek 0
          lift $ action hdr h
  
      ExceptT $ return r
  