{-# LANGUAGE LambdaCase      #-}
-- |
-- Module      : Database.RocksDB.Base
-- Copyright   : (c) 2012-2013 The leveldb-haskell Authors
--               (c) 2014-2020 The rocksdb-haskell Authors
-- License     : BSD3
-- Maintainer  : jprupp@protonmail.ch
-- Stability   : experimental
-- Portability : non-portable
--
-- RocksDB Haskell binding.
--
-- The API closely follows the C-API of RocksDB.

module Database.RocksDB.Base
    ( -- * Exported Types
      DB (..)
    , BatchOp (..)
    , Range

    -- * Options
    , Config (..)

    -- * Basic Database Manipulations
    , withDB
    , put
    , delete
    , write
    , get
    , withSnapshot

    -- * Filter Policy / Bloom Filter
    , BloomFilter
    , withBloomFilter

    -- * Administrative Functions
    , Property (..), getProperty
    , destroy
    , repair
    , approximateSize

    -- * Iteration
    , module Database.RocksDB.Iterator
    ) where

import           Control.Monad             ((>=>))
import           Data.ByteString           (ByteString)
import qualified Data.ByteString           as BS
import           Data.ByteString.Internal  (ByteString (..))
import qualified Data.ByteString.Unsafe    as BU
import           Database.RocksDB.C
import           Database.RocksDB.Internal
import           Database.RocksDB.Iterator
import           Foreign
import           Foreign.C.String          (CString, withCString)
import qualified GHC.Foreign               as GHC
import qualified GHC.IO.Encoding           as GHC
import           System.Directory          (createDirectoryIfMissing)
import           UnliftIO

-- | Properties exposed by RocksDB
data Property = NumFilesAtLevel Int | Stats | SSTables
    deriving (Eq, Show)

data BatchOp = Put !ByteString !ByteString
             | Del !ByteString
             | PutCF !ColumnFamily !ByteString !ByteString
             | DelCF !ColumnFamily !ByteString
             deriving (Eq, Show)

-- | Create a 'BloomFilter'
withBloomFilter :: MonadUnliftIO m => Int -> (BloomFilter -> m a) -> m a
withBloomFilter i =
    bracket create_bloom_filter destroy_bloom_filter
  where
    destroy_bloom_filter = liftIO . c_rocksdb_filterpolicy_destroy
    create_bloom_filter = liftIO $
        c_rocksdb_filterpolicy_create_bloom (intToCInt i)

-- | Open a database.
--
-- The returned handle should be released with 'close'.
withDB :: MonadUnliftIO m => FilePath -> Config -> (DB -> m a) -> m a
withDB path config f =
    withOptions config $ \opts_ptr ->
    withReadOpts Nothing $ \read_opts ->
    withWriteOpts $ \write_opts ->
    bracket (create_db opts_ptr read_opts write_opts) destroy_db f
  where
    destroy_db db = liftIO $
        c_rocksdb_close $ rocksDB db
    create_db opts_ptr read_opts write_opts = liftIO $ do
        createDirectoryIfMissing True path
        withFilePath path $ \path_ptr -> do
            db_ptr <- throwIfErr "open" $
                c_rocksdb_open opts_ptr path_ptr
            return $ DB db_ptr [] read_opts write_opts

-- | Run an action with a 'Snapshot' of the database.
withSnapshot :: MonadUnliftIO m => DB -> (DB -> m a) -> m a
withSnapshot db@DB{rocksDB = db_ptr} f =
    bracket create_snapshot release_snapshot (f . fst)
  where
    release_snapshot = liftIO . c_rocksdb_release_snapshot db_ptr . snd
    create_snapshot = liftIO $ do
        snap_ptr <- c_rocksdb_create_snapshot db_ptr
        withReadOpts (Just snap_ptr) $ \read_opts ->
            return (db{readOpts = read_opts}, snap_ptr)

-- | Get a DB property.
getProperty :: MonadIO m => DB -> Property -> m (Maybe ByteString)
getProperty DB{rocksDB = db_ptr} p = liftIO $
    withCString (prop p) $
    c_rocksdb_property_value db_ptr >=> \case
    val_ptr | val_ptr == nullPtr -> return Nothing
            | otherwise -> do
                  res <- Just <$> BS.packCString val_ptr
                  freeCString val_ptr
                  return res
  where
    prop (NumFilesAtLevel i) = "rocksdb.num-files-at-level" ++ show i
    prop Stats               = "rocksdb.stats"
    prop SSTables            = "rocksdb.sstables"

-- | Destroy the given RocksDB database.
destroy :: MonadIO m => FilePath -> Options -> m ()
destroy path opts_ptr = liftIO $
    withFilePath path $ \path_ptr ->
    throwIfErr "destroy" $ c_rocksdb_destroy_db opts_ptr path_ptr

-- | Repair the given RocksDB database.
repair :: MonadIO m => FilePath -> Options -> m ()
repair path opts_ptr = liftIO $
    withFilePath path $ \path_ptr ->
    throwIfErr "repair" $ c_rocksdb_repair_db opts_ptr path_ptr


-- TODO: support [Range], like C API does
type Range  = (ByteString, ByteString)

-- | Inspect the approximate sizes of the different levels.
approximateSize :: MonadIO m => DB -> Range -> m Int64
approximateSize DB{rocksDB = db_ptr} (from, to) = liftIO $
    BU.unsafeUseAsCStringLen from $ \(from_ptr, flen) ->
    BU.unsafeUseAsCStringLen to   $ \(to_ptr, tlen)   ->
    withArray [from_ptr]          $ \from_ptrs        ->
    withArray [intToCSize flen]   $ \flen_ptrs        ->
    withArray [to_ptr]            $ \to_ptrs          ->
    withArray [intToCSize tlen]   $ \tlen_ptrs        ->
    allocaArray 1                 $ \size_ptrs        -> do
        c_rocksdb_approximate_sizes db_ptr 1
                                    from_ptrs flen_ptrs
                                    to_ptrs tlen_ptrs
                                    size_ptrs
        fmap head $ peekArray 1 size_ptrs >>= mapM toInt64

    where
        toInt64 = return . fromIntegral

-- | Write a key/value pair.
put :: MonadIO m => DB -> ByteString -> ByteString -> m ()
put DB{rocksDB = db_ptr, writeOpts = write_opts} key value = liftIO $
    BU.unsafeUseAsCStringLen key   $ \(key_ptr, klen) ->
    BU.unsafeUseAsCStringLen value $ \(val_ptr, vlen) ->
        throwIfErr "put"
            $ c_rocksdb_put db_ptr write_opts
                            key_ptr (intToCSize klen)
                            val_ptr (intToCSize vlen)

-- | Read a value by key.
get :: MonadIO m => DB -> ByteString -> m (Maybe ByteString)
get DB{rocksDB = db_ptr, readOpts = read_opts} key = liftIO $
    BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
    alloca                       $ \vlen_ptr -> do
        val_ptr <- throwIfErr "get" $
            c_rocksdb_get db_ptr read_opts key_ptr (intToCSize klen) vlen_ptr
        vlen <- peek vlen_ptr
        if val_ptr == nullPtr
            then return Nothing
            else do
                res' <- Just <$> BS.packCStringLen (val_ptr, cSizeToInt vlen)
                freeCString val_ptr
                return res'

-- | Delete a key/value pair.
delete :: MonadIO m => DB -> ByteString -> m ()
delete DB{rocksDB = db_ptr, writeOpts = write_opts} key = liftIO $
    BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
        throwIfErr "delete"
            $ c_rocksdb_delete db_ptr write_opts key_ptr (intToCSize klen)

-- | Perform a batch mutation.
write :: MonadIO m => DB -> [BatchOp] -> m ()
write DB{rocksDB = db_ptr, writeOpts = write_opts} batch = liftIO $
    bracket
    c_rocksdb_writebatch_create
    c_rocksdb_writebatch_destroy $ \batch_ptr -> do

    mapM_ (batchAdd batch_ptr) batch

    throwIfErr "write" $ c_rocksdb_write db_ptr write_opts batch_ptr

    -- ensure @ByteString@s (and respective shared @CStringLen@s) aren't GC'ed
    -- until here
    mapM_ (liftIO . touch) batch

    where
        batchAdd batch_ptr (Put key val) =
            BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
            BU.unsafeUseAsCStringLen val $ \(val_ptr, vlen) ->
                c_rocksdb_writebatch_put
                batch_ptr
                key_ptr (intToCSize klen)
                val_ptr (intToCSize vlen)

        batchAdd batch_ptr (PutCF cf_ptr key val) =
            BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
            BU.unsafeUseAsCStringLen val $ \(val_ptr, vlen) ->
                c_rocksdb_writebatch_put_cf
                batch_ptr
                cf_ptr
                key_ptr (intToCSize klen)
                val_ptr (intToCSize vlen)

        batchAdd batch_ptr (Del key) =
            BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
                c_rocksdb_writebatch_delete
                batch_ptr
                key_ptr (intToCSize klen)

        batchAdd batch_ptr (DelCF cf_ptr key) =
            BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
                c_rocksdb_writebatch_delete_cf
                batch_ptr
                cf_ptr
                key_ptr (intToCSize klen)


        touch (Put (PS p _ _) (PS p' _ _)) = do
            touchForeignPtr p
            touchForeignPtr p'

        touch (PutCF _ (PS p _ _) (PS p' _ _)) = do
            touchForeignPtr p
            touchForeignPtr p'

        touch (Del (PS p _ _)) = touchForeignPtr p

        touch (DelCF _ (PS p _ _)) = touchForeignPtr p

-- | Marshal a 'FilePath' (Haskell string) into a `NUL` terminated C string using
-- temporary storage.
withFilePath :: FilePath -> (CString -> IO a) -> IO a
withFilePath = GHC.withCString GHC.utf8