{-# 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
    , ColumnFamily

    -- * Options
    , Config (..)

    -- * Basic Database Manipulations
    , withDB
    , withDBCF
    , put
    , putCF
    , delete
    , deleteCF
    , write
    , get
    , getCF
    , withSnapshot

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

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

import           Control.Monad             (when, (>=>), forM)
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           UnliftIO
import           UnliftIO.Directory
import           UnliftIO.Foreign

-- | 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)

-- | 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 = do
        when (createIfMissing config) $
            createDirectoryIfMissing True path
        withCString path $ \path_ptr -> do
            db_ptr <- liftIO . throwIfErr "open" $
                c_rocksdb_open opts_ptr path_ptr
            return DB { rocksDB = db_ptr
                      , columnFamilies = []
                      , readOpts = read_opts
                      , writeOpts = write_opts
                      }

withDBCF :: MonadUnliftIO m
         => FilePath
         -> Config
         -> [(String, Config)]
         -> (DB -> m a)
         -> m a
withDBCF path config cf_cfgs f =
    withOptions config $ \opts_ptr ->
    withOptionsCF (map snd cf_cfgs) $ \cf_opts ->
    withReadOpts Nothing $ \read_opts ->
    withWriteOpts $ \write_opts ->
    withStrings (map fst cf_cfgs) $ \cf_names ->
    allocaArray (length cf_cfgs + 1) $ \cf_names_array ->
    allocaArray (length cf_cfgs + 1) $ \cf_opts_array ->
    allocaArray (length cf_cfgs + 1) $ \cf_ptrs_array ->
        bracket ( create_db opts_ptr
                            cf_names
                            cf_names_array
                            cf_opts
                            cf_opts_array
                            read_opts
                            cf_ptrs_array
                            write_opts
                ) destroy_db f
  where
    create_new cf_names cf_opts opts_ptr read_opts write_opts = do
        createDirectoryIfMissing True path
        withCString path $ \path_ptr -> do
            db_ptr <- liftIO . throwIfErr "open" $
                c_rocksdb_open opts_ptr path_ptr
            cfs <- forM (zip cf_names cf_opts) $ \(n, o) ->
                throwIfErr "create_column_family" $
                c_rocksdb_create_column_family
                db_ptr o n
            return DB { rocksDB = db_ptr
                      , columnFamilies = cfs
                      , readOpts = read_opts
                      , writeOpts = write_opts
                      }
    destroy_db db = liftIO $ do
        mapM_ c_rocksdb_column_family_handle_destroy (columnFamilies db)
        c_rocksdb_close $ rocksDB db
    create_db opts_ptr
              cf_names
              cf_names_array
              cf_opts
              cf_opts_array
              read_opts
              cf_ptrs_array
              write_opts = liftIO $ do
        when (createIfMissing config) $
            createDirectoryIfMissing True path
        null <$> listDirectory path >>= \case
            True -> create_new cf_names cf_opts opts_ptr read_opts write_opts
            False -> withCString path $ \path_ptr ->
                withCString "default" $ \cf_deflt_name -> do
                    pokeArray cf_names_array (cf_deflt_name : cf_names)
                    pokeArray cf_opts_array (opts_ptr : cf_opts)
                    db_ptr <- throwIfErr "open" $
                        c_rocksdb_open_column_families
                        opts_ptr
                        path_ptr
                        (intToCInt (length cf_cfgs + 1))
                        cf_names_array
                        cf_opts_array
                        cf_ptrs_array
                    cfs <- peekArray (length cf_cfgs + 1) cf_ptrs_array
                    return DB { rocksDB = db_ptr
                              , columnFamilies = tail cfs
                              , readOpts = read_opts
                              , writeOpts = write_opts
                              }

-- | Run an action with a snapshot of the database.
-- The 'DB' object is not valid after the action ends.
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 $
    withCString 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 $
    withCString 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 = putCommon db Nothing

putCF :: MonadIO m => DB -> ColumnFamily -> ByteString -> ByteString -> m ()
putCF db cf = putCommon db (Just cf)

putCommon :: MonadIO m => DB -> Maybe ColumnFamily -> ByteString -> ByteString -> m ()
putCommon DB{rocksDB = db_ptr, writeOpts = write_opts} mcf key value = liftIO $
    BU.unsafeUseAsCStringLen key   $ \(key_ptr, klen) ->
    BU.unsafeUseAsCStringLen value $ \(val_ptr, vlen) ->
        throwIfErr "put" $ case mcf of
            Just cf -> c_rocksdb_put_cf
                      db_ptr write_opts cf
                      key_ptr (intToCSize klen)
                      val_ptr (intToCSize vlen)
            Nothing -> 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 = getCommon db Nothing

getCF :: MonadIO m => DB -> ColumnFamily -> ByteString -> m (Maybe ByteString)
getCF db cf = getCommon db (Just cf)

getCommon :: MonadIO m => DB -> Maybe ColumnFamily -> ByteString -> m (Maybe ByteString)
getCommon DB{rocksDB = db_ptr, readOpts = read_opts} mcf key = liftIO $
    BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
    alloca                       $ \vlen_ptr -> do
        val_ptr <- throwIfErr "get" $
            case mcf of
                Just cf -> c_rocksdb_get_cf
                           db_ptr read_opts cf
                           key_ptr (intToCSize klen) vlen_ptr
                Nothing -> 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 :: MonadIO m => DB -> ByteString -> m ()
delete db = deleteCommon db Nothing

deleteCF :: MonadIO m => DB -> ColumnFamily -> ByteString -> m ()
deleteCF db cf = deleteCommon db (Just cf)

-- | Delete a key/value pair.
deleteCommon :: MonadIO m => DB -> Maybe ColumnFamily -> ByteString -> m ()
deleteCommon DB{rocksDB = db_ptr, writeOpts = write_opts} mcf key = liftIO $
    BU.unsafeUseAsCStringLen key $ \(key_ptr, klen) ->
    throwIfErr "delete" $ case mcf of
    Just cf -> c_rocksdb_delete_cf db_ptr write_opts cf key_ptr (intToCSize klen)
    Nothing -> 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

withStrings :: MonadUnliftIO m => [String] -> ([CString] -> m a) -> m a
withStrings ss f =
    go [] ss
  where
    go acc [] = f (reverse acc)
    go acc (x:xs) = withCString x $ \p -> go (p:acc) xs