{-# LANGUAGE DataKinds #-}

module Jerimum.Tests.Unit.Storage.LMDB.SegmentTest
  ( tests
  ) where

import           Control.Concurrent
import           Control.Exception
import           Control.Monad
import qualified Data.ByteString                        as B
import           Data.Either
import           Data.Function
import           Data.List
import           Jerimum.Storage.LMDB.Segment
import           Jerimum.Tests.Unit.Storage.LMDB.Helper
import           System.IO.Temp
import           Test.QuickCheck.Monadic                as QM
import           Test.Tasty
import           Test.Tasty.HUnit
import           Test.Tasty.QuickCheck                  as Q

genKeyVals :: Monad m => PropertyM m [(Key, Value)]
genKeyVals =
  sortBy (compare `on` fst) .
  nubBy ((==) `on` fst) . map (\(k, v) -> (B.pack (0 : k), B.pack v)) <$>
  pick arbitrary

canFetchAnyKeyWritten :: IO (Handle 'RW) -> Property
canFetchAnyKeyWritten getHandle =
  monadicIO $ do
    key <- B.pack <$> pick (arbitrary `suchThat` (not . null))
    val <- B.pack <$> pick arbitrary
    handle <- run getHandle
    writeOk <-
      run $ wtransaction handle $ \txn -> TxCommit <$> putKey txn key val
    readVal <- run $ rtransaction handle (`getKey` key)
    QM.assert (Right True == (writeOk :: Either () Bool) && Just val == readVal)

scanKeysInOrder :: IO (Handle 'RW) -> Property
scanKeysInOrder getHandle =
  monadicIO $ do
    handle <- run getHandle
    keyvals <- genKeyVals
    _ <-
      run $
      wtransaction handle $ \txn ->
        TxCommit <$> mapM_ (uncurry (putKey txn)) keyvals
    allKeys <-
      run $ rtransaction handle (\txn -> scan txn Nothing (length keyvals))
    someKeys <-
      run $
      rtransaction handle $ \txn ->
        mapM
          ((\k -> scan txn (Just k) (length keyvals)) . fst . head)
          (init $ tails keyvals)
    QM.assert (keyvals == allKeys && someKeys == init (tails keyvals))

orderedWritesTest :: TestTree
orderedWritesTest =
  testCase "ordered writes" $
  withSegment $ \handle -> do
    _ <- wtransaction handle (\txn -> TxCommit <$> putKeys txn ["0", "1", "2"])
    keys <- map fst <$> rtransaction handle (\txn -> scan txn Nothing 3)
    keys @?= ["0", "1", "2"]

outOfOrderWritesTest :: TestTree
outOfOrderWritesTest =
  testCase "out of order writes" $
  withSegment $ \handle -> do
    _ <-
      wtransaction
        handle
        (\txn -> TxCommit <$> putKeys txn (reverse ["0", "1", "2"]))
    keys <- map fst <$> rtransaction handle (\txn -> scan txn Nothing 3)
    keys @?= ["0", "1", "2"]

sizeLimitTest :: TestTree
sizeLimitTest =
  testCase "size limit" $
  withSegment $ \handle -> do
    let largeValue = B.replicate (4 * 1024 * 1024) 0x7a
    ans <-
      tryJust (guard . isMapFullError) $
      wtransaction handle (\txn -> TxCommit <$> putKey txn "0" largeValue)
    isLeft ans @? "no exception was thrown"

sizeLimitOverwriteTest :: TestTree
sizeLimitOverwriteTest =
  testCase "size limit | overwrite" $
  withSegment $ \handle -> do
    let largeValue = B.replicate (1024 * 1024) 0x7a
    _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "")
    ans <-
      wtransaction handle (\txn -> TxCommit <$> putNewKey txn "0" largeValue)
    ans @?= (Right False :: Either () Bool)

putNewTest :: TestTree
putNewTest =
  testCase "putNew" $
  withSegment $ \handle -> do
    _ <- wtransaction handle (\txn -> TxCommit <$> putKeys txn ["0", "1", "2"])
    ans <-
      wtransaction handle (\txn -> TxCommit <$> putNewKeys txn ["0", "1", "2"])
    ans @?= (Right (replicate 3 False) :: Either () [Bool])

txRollbackAndCommitTest :: TestTree
txRollbackAndCommitTest =
  testCase "TxRollback and TxCommit" $
  withSegment $ \handle -> do
    _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "foobar")
    ans <- rtransaction handle (`getKey` "0")
    ans @?= Just "foobar"
    _ <- wtransaction handle (\txn -> TxRollback <$> putKey txn "0" "foobaz")
    ans <- rtransaction handle (`getKey` "0")
    ans @?= Just "foobar"
    _ <- wtransaction handle (\txn -> TxRollback <$> putKey txn "1" "foobar")
    ans <- rtransaction handle (`getKey` "1")
    ans @?= Nothing

tests :: IO (Handle 'RW) -> TestTree
tests getHandle =
  testGroup
    "Storage.LMDB.Segment"
    [ testGroup
        "writes"
        [ orderedWritesTest
        , outOfOrderWritesTest
        , sizeLimitTest
        , sizeLimitOverwriteTest
        , putNewTest
        ]
    , testProperty "can fetch any key written" (canFetchAnyKeyWritten getHandle)
    , testProperty "scan keys in order" (scanKeysInOrder getHandle)
    , testGroup "transactions" [txRollbackAndCommitTest]
    , testGroup
        "reads"
        [ testCase "scan from missing key" $
          withSegment $ \handle -> do
            _ <-
              wtransaction
                handle
                (\txn -> TxCommit <$> putKeys txn ["5", "6", "7"])
            keys <-
              map fst <$> rtransaction handle (\txn -> scan txn (Just "0") 3)
            keys @?= ["5", "6", "7"]
        , testCase "scan from existing key" $
          withSegment $ \handle -> do
            _ <-
              wtransaction
                handle
                (\txn -> TxCommit <$> putKeys txn ["0", "1", "2"])
            keys <-
              map fst <$> rtransaction handle (\txn -> scan txn (Just "1") 3)
            keys @?= ["1", "2"]
        , testCase "bounds on an empty db" $
          withSegment $ \handle -> do
            ans <- rtransaction handle bounds
            ans @?= Nothing
        , testCase "bounds after single write" $
          withSegment $ \handle -> do
            _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "")
            ans <- rtransaction handle bounds
            ans @?= Just ("0", "0")
        , testCase "bounds after series of writes" $
          withSegment $ \handle -> do
            _ <-
              wtransaction
                handle
                (\txn -> TxCommit <$> putKeys txn ["0", "1", "2"])
            ans <- rtransaction handle bounds
            ans @?= Just ("0", "2")
        ]
    , testGroup
        "information"
        [ testCase "getPath" $
          withSystemTempDirectory "nws-storage-unit" $ \tmpdir ->
            bracket (open tmpdir (1024 * 1024) rwMode) close $ \handle -> do
              path <- getPath handle
              path @?= tmpdir
        ]
    , testGroup
        "concurrency"
        [ testCase "putKey" $
          withSegment $ \handle -> do
            mvars <- mapM (\k -> putKeyAsync handle k "") ["0", "1", "2"]
            mapM_ takeMVar mvars
            keys <- map fst <$> rtransaction handle (\txn -> scan txn Nothing 3)
            keys @?= ["0", "1", "2"]
        , testCase "putNewKey | async" $
          withSegment $ \handle -> do
            mvars <- replicateM 10 (putNewKeyAsync handle "0" "")
            ans <- mapM takeMVar mvars
            length (filter (== Just True) ans) @?= 1
            length (filter (== Just False) ans) @?= 9
        , testCase "scan" $
          withSegment $ \handle -> do
            _ <-
              wtransaction
                handle
                (\txn -> TxCommit <$> putKeys txn ["0", "1", "2"])
            mvars <- mapM (const $ getKeysAsync handle) ([0 .. 9] :: [Int])
            keys <- mapM takeMVar mvars
            keys @?=
              [[("0", ""), ("1", ""), ("2", "")] | _ <- [0 .. 9] :: [Int]]
        , testCase "write does not block reads" $
          withSegment $ \handle -> do
            mvar0 <- newEmptyMVar
            mvar1 <- newEmptyMVar
            _ <-
              forkFinally
                (takeMVar mvar0 >> rtransaction handle (`getKey` "0"))
                (\(Right value) -> putMVar mvar1 value)
            value <-
              runBeforeWrite handle "0" "foobar" $ do
                putMVar mvar0 ()
                takeMVar mvar1
            value @?= Nothing
        ]
    , testGroup
        "consistency"
        [ testCase "read-after-create" $
          withSegment $ \handle -> do
            _ <-
              wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "foobar")
            value <- rtransaction handle (`getKey` "0")
            value @?= Just "foobar"
        , testCase "read-after-create (in-transaction)" $ do
            value <-
              withSegment $ \handle ->
                wtransaction handle $ \txn -> do
                  _ <- putKey txn "0" "foobar"
                  TxCommit <$> getKey txn "0"
            value @?= (Right (Just "foobar") :: Either () (Maybe Value))
        , testCase "read-after-update" $
          withSegment $ \handle -> do
            _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "foo")
            _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "bar")
            value <- rtransaction handle (`getKey` "0")
            value @?= Just "bar"
        , testCase "read-after-update (in-transaction)" $
          withSegment $ \handle -> do
            value <-
              wtransaction handle $ \txn -> do
                _ <- putKey txn "0" "foo"
                _ <- putKey txn "0" "bar"
                TxCommit <$> getKey txn "0"
            value @?= (Right (Just "bar") :: Either () (Maybe Value))
        , testCase "read-after-delete" $
          withSegment $ \handle -> do
            _ <- wtransaction handle (\txn -> TxCommit <$> putKey txn "0" "foo")
            _ <- wtransaction handle (\txn -> TxCommit <$> delKey txn "0")
            value <- rtransaction handle (`getKey` "0")
            value @?= Nothing
        , testCase "read-after-delete (in-transaction)" $
          withSegment $ \handle -> do
            value <-
              wtransaction handle $ \txn -> do
                _ <- putKey txn "0" "foo"
                _ <- delKey txn "0"
                TxCommit <$> getKey txn "0"
            value @?= (Right Nothing :: Either () (Maybe Value))
        ]
    ]
