{-# 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
    , snapshot
    , createSnapshot
    , releaseSnapshot

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

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

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

-- | 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 f =
    bracket (createSnapshot db) releaseSnapshot (f . fst)

-- | The 'DB' snapshot is not valid outside of 'MonadResource'.
snapshot :: (MonadIO m, MonadResource m) => DB -> m DB
snapshot db =
    fst . snd <$> allocate (createSnapshot db) releaseSnapshot

-- | Manually create an unmanaged snapshot.
createSnapshot :: MonadIO m => DB -> m (DB, Snapshot)
createSnapshot db@DB{rocksDB = db_ptr} = liftIO $ do
    snap_ptr <- c_rocksdb_create_snapshot db_ptr
    withReadOpts (Just snap_ptr) $ \read_opts ->
        return (db{readOpts = read_opts}, snap_ptr)

-- | Function to release an unmanaged snapshot.
releaseSnapshot :: MonadIO m => (DB, Snapshot) -> m ()
releaseSnapshot (DB{rocksDB = db_ptr}, snap_ptr) =
    liftIO $ c_rocksdb_release_snapshot db_ptr 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