import System.Directory
import Database.Berkeley.Db
import Test.HUnit
import Control.Exception.Extensible
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

withCase1 action = do
    cleanDb $ \dbenv db -> do
        dbEnv_withTxn [] [] dbenv Nothing $ \txn -> do
            db_put [] db (Just txn) (pack "car")   (pack "triumph")
            db_put [] db (Just txn) (pack "fruit") (pack "banana")
            db_put [] db (Just txn) (pack "car")   (pack "honda")
            db_put [] db (Just txn) (pack "tree")  (pack "poplar")
        action dbenv db

test_putget = do
    withCase1 $ \dbenv db -> do
        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
            locker <- dbTxn_id txn
            dbEnv_withLock [] DB_LOCK_WRITE (pack "my_lock") dbenv locker $ do
                db_put [] db (Just txn) (pack "fruit") (pack "durian")

test_dbCursor_get =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur -> do
        assertEqual "dbCursor_get_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_get_fruit" (Just (pack "fruit", pack "banana")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_get_tree" (Just (pack "tree", pack "poplar")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_get_Nothing" Nothing =<< dbCursor_get [DB_NEXT] cur

test_dbCursor_put1 =
    cleanDb $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn -> do
        db_withCursor [] db (Just txn) $ \cur -> do
            dbCursor_put [DB_KEYFIRST] cur (pack "dog") (pack "Xanthe")
            dbCursor_put [DB_KEYFIRST] cur (pack "cat") (pack "Jasper")
        db_withCursor [] db (Just txn) $ \cur -> do
            assertEqual "dbCursor_put1_cat" (Just (pack "cat", pack "Jasper")) =<< dbCursor_get [DB_NEXT] cur
            assertEqual "dbCursor_put1_dog" (Just (pack "dog", pack "Xanthe")) =<< dbCursor_get [DB_NEXT] cur

test_dbCursor_put2 =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur -> do
        assertEqual "dbCursor_put2_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_put2_fruit" (Just (pack "fruit", pack "banana")) =<< dbCursor_get [DB_NEXT] cur
        dbCursor_put [DB_CURRENT] cur (pack "") (pack "orange")
        assertEqual "dbCursor_put2_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_FIRST] cur
        assertEqual "dbCursor_put2_fruit" (Just (pack "fruit", pack "orange")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_put2_tree" (Just (pack "tree", pack "poplar")) =<< dbCursor_get [DB_NEXT] cur

test_dbCursor_set =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur -> do
        assertEqual "dbCursor_set_tree" (Just $ pack "poplar") =<< dbCursor_set [] cur (pack "tree")
        assertEqual "dbCursor_set_Nothing" Nothing =<< dbCursor_get [DB_NEXT] cur

test_dbCursor_dup =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur1 -> do
        assertEqual "dbCursor_dup_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_NEXT] cur1
        dbCursor_withCursor [DB_POSITION] cur1 $ \cur2 -> do
            assertEqual "dbCursor_dup_2_fruit" (Just (pack "fruit", pack "banana")) =<< dbCursor_get [DB_NEXT] cur2
            assertEqual "dbCursor_dup_2_tree" (Just (pack "tree", pack "poplar")) =<< dbCursor_get [DB_NEXT] cur2
        assertEqual "dbCursor_dup_fruit" (Just (pack "fruit", pack "banana")) =<< dbCursor_get [DB_NEXT] cur1
        assertEqual "dbCursor_dup_tree" (Just (pack "tree", pack "poplar")) =<< dbCursor_get [DB_NEXT] cur1
        assertEqual "dbCursor_dup_Nothing" Nothing =<< dbCursor_get [DB_NEXT] cur1

test_dbCursor_del =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur -> do
        assertEqual "dbCursor_del_1_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_del_1_fruit" (Just (pack "fruit", pack "banana")) =<< dbCursor_get [DB_NEXT] cur
        dbCursor_del [] cur
        assertEqual "dbCursor_del_2_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_FIRST] cur
        assertEqual "dbCursor_del_2_tree" (Just (pack "tree", pack "poplar")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_del_2_Nothing" Nothing =<< dbCursor_get [DB_NEXT] cur

test_dbCursor_count =
    withCase1 $ \dbenv db ->
    dbEnv_withTxn [] [] dbenv Nothing $ \txn ->
    db_withCursor [] db (Just txn) $ \cur -> do
        assertEqual "dbCursor_count_car" (Just (pack "car", pack "honda")) =<< dbCursor_get [DB_NEXT] cur
        assertEqual "dbCursor_count_val" 1 =<< dbCursor_count [] cur

main = do
    runTestTT $ TestList [
            TestCase test_exception,
            TestCase test_putget,
            TestCase test_txn,
            TestCase test_afterclose,
            TestCase test_lock,
            TestCase test_dbCursor_get,
            TestCase test_dbCursor_put1,
            TestCase test_dbCursor_put2,
            TestCase test_dbCursor_set,
            TestCase test_dbCursor_dup,
            TestCase test_dbCursor_del,
            TestCase test_dbCursor_count
        ]