import System.Directory
import Database.Berkeley.Db
import Test.HUnit
import Control.Exception
import Control.Monad
import qualified Data.ByteString as B
import Data.Char
import Prelude hiding (catch)


pack :: String -> B.ByteString
pack = B.pack . map (fromIntegral . ord)

unpack :: B.ByteString -> String
unpack = map (chr . fromIntegral) . B.unpack 

cleanDir :: IO () -> IO ()
cleanDir code = do
    original <- getCurrentDirectory
    removeDirectoryRecursive "test" `catch` ((\exc -> return ())::SomeException -> IO ())
    createDirectory "test"
    setCurrentDirectory "test"
    code
    setCurrentDirectory original

cleanEnv :: (DbEnv -> IO ()) -> IO ()
cleanEnv code = do
    cleanDir $
        bracket
            (do
                dbenv <- dbEnv_create []
                dbEnv_open [DB_CREATE,DB_INIT_LOCK,DB_INIT_LOG,DB_INIT_MPOOL,
                    DB_INIT_TXN,DB_THREAD] 0 dbenv "."
                return dbenv)
            (dbEnv_close [])
            code

cleanDb :: (DbEnv -> Db -> IO ()) -> IO ()
cleanDb code =
    cleanEnv $ \dbenv -> do
        bracket
            (do
                db <- db_create [] dbenv
                db_open [DB_CREATE,DB_THREAD,DB_AUTO_COMMIT] DB_BTREE 0 db Nothing "test.db" Nothing
                return db)
            (db_close [])
            (code dbenv)

test_exception = do
    cleanEnv $ \dbenv -> do
        answer <-
            catch
                (do
                    bracket
                        (db_create [] dbenv)
                        (db_close [])
                        (\db -> db_open [] DB_HASH 0 db Nothing "not-there.db" Nothing)
                    return Nothing)
                (\exc -> do
                        case fromException exc of
                            Just (DbException _ code) -> return $ Just code
                            Nothing                   -> throwIO exc
                    )
        assertEqual "exception" (Just (SYSTEM_ERROR 2)) answer  -- 2 == ENOENT

test_putget = do
    cleanDb $ \dbenv db -> do
        db_put [] db Nothing (pack "car")   (pack "triumph")
        db_put [] db Nothing (pack "fruit") (pack "banana")
        db_put [] db Nothing (pack "car")   (pack "honda")
        db_put [] db Nothing (pack "tree")  (pack "poplar")
        assertEqual "putget_fruit" (Just $ pack "banana") =<< db_get [] db Nothing (pack "fruit")
        assertEqual "putget_car"   (Just $ pack "honda")  =<< db_get [] db Nothing (pack "car")
        assertEqual "putget_tree"  (Just $ pack "poplar") =<< db_get [] db Nothing (pack "tree")
        assertEqual "putget_pet"   Nothing                =<< db_get [] db Nothing (pack "pet")

test_txn = do
    cleanDb $ \dbenv db -> do
        dbEnv_withTxn [] [] dbenv Nothing $ \txn -> do
            db_put [] db (Just txn) (pack "fruit") (pack "apricot")
        assertEqual "txn_1" (Just $ pack "apricot") =<< db_get [] db Nothing (pack "fruit")
          -- Check that a rolled back exception doesn't modify the database
        catch
            (dbEnv_withTxn [] [] dbenv Nothing $ \txn -> do
                db_put [] db (Just txn) (pack "fruit") (pack "pomegranate")
                throwIO $ DbException "" (SYSTEM_ERROR 0))
            ((\exc -> return ()) :: SomeException -> IO ())
        assertEqual "txn_1" (Just $ pack "apricot") =<< db_get [] db Nothing (pack "fruit")

test_afterclose = do
    cleanDir $ do
        mDbErr <- catch (do
            dbenv <- dbEnv_create []
            dbEnv_open [DB_CREATE,DB_INIT_LOCK,DB_INIT_LOG,DB_INIT_MPOOL,
                        DB_INIT_TXN,DB_THREAD] 0 dbenv "."
            dbEnv_close [] dbenv
            db <- db_create [] dbenv
            db_open [DB_CREATE,DB_THREAD,DB_AUTO_COMMIT] DB_BTREE 0 db Nothing "test.db" Nothing
            db_close [] db
            return Nothing)
            (\exc ->
                case fromException exc of
                    Just (DbException _ code) -> return $ Just code
                    Nothing                   -> throwIO exc)
        assertEqual "afterclose_1" (Just DB_ACCESSED_DB_ENV_AFTER_CLOSE) mDbErr

    cleanEnv $ \dbenv -> do
        mDbErr <- catch (do
            db <- db_create [] dbenv
            db_open [DB_CREATE,DB_THREAD,DB_AUTO_COMMIT] DB_BTREE 0 db Nothing "test.db" Nothing
            db_close [] db
            db_put [] db Nothing (pack "fruit") (pack "guava")
            return Nothing)
            (\exc ->
                case fromException exc of
                    Just (DbException _ code) -> return $ Just code
                    Nothing                   -> throwIO exc)
        assertEqual "afterclose_2" (Just DB_ACCESSED_DB_AFTER_CLOSE) mDbErr
        
    cleanDb $ \dbenv db -> do
        mDbErr <- catch (do
            txn <- dbEnv_txn_begin [] dbenv Nothing
            dbTxn_abort txn
            db_put [] db (Just txn) (pack "fruit") (pack "melon")
            return Nothing)
            (\exc ->
                case fromException exc of
                    Just (DbException _ code) -> return $ Just code
                    Nothing                   -> throwIO exc)
        assertEqual "afterclose_3" (Just DB_ACCESSED_DB_TXN_AFTER_CLOSE) mDbErr

test_lock = do
    cleanDb $ \dbenv db -> do
        -- This doesn't test very much: Just that locking doesn't cause any
        -- runtime errors, hangs or crashes.
        dbEnv_withTxn [] [] dbenv Nothing $ \txn -> do
            dbEnv_withLock [] DB_LOCK_WRITE (pack "my_lock") dbenv (dbTxn_id txn) $ do
                db_put [] db (Just txn) (pack "fruit") (pack "durian")

main = do
    runTestTT $ TestList [
            TestCase test_exception,
            TestCase test_putget,
            TestCase test_txn,
            TestCase test_afterclose,
            TestCase test_lock
        ]