{-# LANGUAGE RankNTypes #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE DeriveDataTypeable #-}
{-# LANGUAGE RecordWildCards #-}

{-# OPTIONS -Wall #-}
{-# OPTIONS -fno-cse #-}

module Main where

import Control.Arrow ((>>>))
import Control.Monad (foldM, when)
import Control.Monad.IO.Class (MonadIO(..))
import Control.Monad.Trans.Resource (runResourceT)
import Data.ByteString.Builder (word32BE, toLazyByteString)
import Data.Default.Class (Default(..))
import Data.Maybe (isJust)
import Data.Word (Word32)
import Data.Data (Data)
import System.Console.CmdArgs (cmdArgs)
import System.IO.Temp (withSystemTempDirectory)

import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as LBS

import qualified Crypto.Hash.BLAKE2.BLAKE2s as Blake2s

import qualified Database.LevelDB as LevelDB
import qualified Database.LMDB.Simple as LMDB

import DFINITY.RadixTree


--------------------------------------------------------------------------------

data Args
  = Args
    { lmdb             :: !Bool
    , lmdb_database    :: !(Maybe FilePath)
    , leveldb          :: !Bool
    , leveldb_database :: !(Maybe FilePath)
    , elements         :: Word32
    }
  deriving (Data)

instance Default Args where
   def = Args False Nothing False Nothing 100000

--------------------------------------------------------------------------------

isDivisibleBy :: (Integral i) => i -> i -> Bool
isDivisibleBy a b = (a `mod` b) == 0

foldN
   :: (RadixDatabase m database, Monad n)
   => Int
   -> (forall a. m a -> n a)
   -> (RadixTree database -> ByteString -> m (RadixTree database))
   -> RadixTree database
   -> [Word32]
   -> n (RadixTree database)
foldN _ _      _      tree [] = pure tree
foldN n commit action tree xs = do
  let toKey = Blake2s.hash 20 mempty . LBS.toStrict . toLazyByteString . word32BE
  tree' <- commit (foldM (\t -> action t . toKey) tree (take n xs)
                   >>= (merkleizeRadixTree >>> fmap snd))
  foldN n commit action tree' (drop n xs)

foldInsert
  :: (RadixDatabase m database, Monad n)
  => (forall a. m a -> n a)
  -> RadixTree database
  -> [Word32]
  -> n (RadixTree database)
foldInsert commit
  = foldN 2000 commit (\db key -> insertRadixTree key key db)

foldDelete
  :: (RadixDatabase m database, Monad n)
  => (forall a. m a -> n a)
  -> RadixTree database
  -> [Word32]
  -> n (RadixTree database)
foldDelete commit
  = foldN 2000 commit (\db key -> deleteRadixTree key db)

--------------------------------------------------------------------------------

withDatabasePath :: Maybe FilePath -> String -> (FilePath -> IO a) -> IO a
withDatabasePath Nothing suffix callback = do
  let name = "dfinity-radix-tree-benchmark-" ++ suffix
  withSystemTempDirectory name callback
withDatabasePath (Just fp) _ callback = do
  callback fp

--------------------------------------------------------------------------------

benchmarkLevelDB :: Args -> IO ()
benchmarkLevelDB args = do
  withDatabasePath (leveldb_database args) "leveldb" $ \dbPath -> do
    let options = def { LevelDB.createIfMissing = True }
    runResourceT $ do
      handle <- LevelDB.open dbPath options
      tree1 <- createRadixTree 262144 2048 Nothing handle
      tree2 <- foldInsert id tree1 [1 .. elements args]
      tree3 <- foldDelete id tree2 [1 .. elements args]
      liftIO $ print $ isEmptyRadixTree tree3

benchmarkLMDB :: Args -> IO ()
benchmarkLMDB args = do
  withDatabasePath (lmdb_database args) "lmdb" $ \dbPath -> do
    let limits = LMDB.defaultLimits { LMDB.mapSize = (1024 * 1024 * 512) }
    env <- LMDB.openReadWriteEnvironment dbPath limits
    db <- LMDB.readOnlyTransaction env (LMDB.getDatabase Nothing)
          :: IO (LMDB.Database ByteString ByteString)
    let runRW = LMDB.readWriteTransaction env
    tree1 <- runRW $ createRadixTree 262144 2048 Nothing db
    tree2 <- foldInsert runRW tree1 [1 .. elements args]
    tree3 <- foldDelete runRW tree2 [1 .. elements args]
    liftIO (print (isEmptyRadixTree tree3))

--------------------------------------------------------------------------------

-- FIXME: maybe move to hs-dfinity-common or whatever; I can see this
--        code getting duplicated

logError, logWarning, logNote :: [String] -> IO ()
logError   msg = putStrLn ("\ESC[31m[ERROR]   " ++ mconcat msg ++ "\ESC[0m")
logWarning msg = putStrLn ("\ESC[33m[WARNING] " ++ mconcat msg ++ "\ESC[0m")
logNote    msg = putStrLn ("\ESC[34m[NOTE]    " ++ mconcat msg ++ "\ESC[0m")

flag :: String -> String
flag s = "\ESC[1m" ++ s ++ "\ESC[21m"

--------------------------------------------------------------------------------

main :: IO ()
main = do
  args <- cmdArgs def
  when (not (lmdb args) && isJust (lmdb_database args)) $ do
    logWarning [ "ignoring argument: "
               , flag "--lmdb-database"
               , " was specified but not "
               , flag "--lmdb"
               ]
  when (not (leveldb args) && isJust (leveldb_database args)) $ do
    logWarning [ "ignoring argument: "
               , flag "--leveldb-database"
               , " was specified but not "
               , flag "--leveldb"
               ]
  when (not (lmdb args) && not (leveldb args)) $ do
    logError [ "neither ", flag "--lmdb", " nor ", flag "--leveldb"
             , " was specified, quitting"
             ]
  when (leveldb args) $ do
    logNote ["benchmarking LevelDB"]
    benchmarkLevelDB args
  when (lmdb args) $ do
    logNote ["benchmarking LMDB"]
    benchmarkLMDB args

--------------------------------------------------------------------------------

-- import qualified Criterion      as C
-- import qualified Criterion.Main as C
--
-- import qualified System.IO
--
-- getRandom :: Int -> IO ByteString
-- getRandom size = do
--   System.IO.withFile "/dev/random" System.IO.ReadMode $ \hdl -> do
--     BS.hGet hdl size
--
-- main :: IO ()
-- main = do
--   let blake2b  = Blake2b.hash  20 mempty
--   let blake2bp = Blake2bp.hash 20 mempty
--   let blake2s  = Blake2s.hash  20 mempty
--   let blake2sp = Blake2sp.hash 20 mempty
--   let sha256   = BSC8.take 20 . SHA256.hash
--
--   let hashB n = C.env (getRandom n)
--                 $ \r -> C.bgroup ("hashing random data of size " ++ show n)
--                         [ C.bench "blake2b-160"     (C.nf blake2b  r)
--                         , C.bench "blake2bp-160"    (C.nf blake2bp r)
--                         , C.bench "blake2s-160"     (C.nf blake2s  r)
--                         , C.bench "blake2sp-160"    (C.nf blake2sp r)
--                         , C.bench "sha256-take-160" (C.nf sha256   r)
--                         ]
--
--   C.defaultMain (map hashB (map (2 ^) ([4 .. 22] :: [Int])))

--------------------------------------------------------------------------------