{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}

module Database.CouchDB.Enumerator.Test.Util (
      CouchT
    , testCouch
    , testCouchCase
    , testCouchProperty
    , isSubmapOf
    , assertStr
    , assertObjMember
    , checkError
    , assertRecvError
    , checkRevision
    , checkLoad
    , clearObject
    , ArbitraryObject(..)
)where

import Control.Applicative
import Control.Monad.IO.Class (MonadIO, liftIO)
import Control.Monad.Trans.Class (lift)
import Control.Monad.Trans.Control (MonadBaseControl)
import Control.Exception.Lifted as E
import Control.Monad
import Control.Monad.Trans.Reader
import qualified Data.Aeson as A
import qualified Data.HashMap.Lazy as M
import Data.Maybe (fromJust)
import Database.CouchDB.Enumerator
import qualified Data.Text as T
import qualified Data.Vector as V

import Test.Framework (Test)
import Test.Framework.Providers.QuickCheck2 (testProperty)
import Test.Framework.Providers.HUnit (testCase)
import Test.QuickCheck
import Test.QuickCheck.Monadic
import Test.HUnit hiding (Test, path)

type CouchT m a = ReaderT CouchConnection m a

testCouch :: CouchT IO a -> IO ()
testCouch c = withCouchConnection "localhost" 5984 "testcouchenum" (runReaderT c) >> return ()

testCouchCase :: String -> CouchT IO a -> Test
testCouchCase s c = testCase s $ testCouch c

testCouchProperty :: (Show a, Arbitrary a) => String -> (Int,Int) -> ([a] -> CouchT IO b) -> Test
testCouchProperty s i f = testProperty s $ monadicIO $ do
    len <- pick $ choose i
    lst <- pick $ vector len
    run $ testCouch $ f lst

-- | Assert that the value is a string, and check that it matches the given string
assertStr :: T.Text -> A.Value -> Assertion
assertStr t (A.String t') = unless (t == t') $ assertFailure $ "strings are not equal. expecting " 
                                                   ++ T.unpack t ++ "  received  " ++ T.unpack t'
assertStr _ _ = assertFailure "expecting a JSON string"

member :: T.Text -> A.Object -> Bool
member k o = M.lookup k o /= Nothing

isSubmapOf :: A.Object -> A.Object -> Bool
isSubmapOf x y = 0 == M.size (M.difference x y)

-- | Assert that the given key exists, and the value matches the given assertion
assertObjMember :: T.Text -> (A.Value -> Assertion) -> A.Object -> Assertion
assertObjMember t f x = do
    assertBool (T.unpack t ++ " is missing") $ member t x
    f $ fromJust $ M.lookup t x

-- | Check an action for a couch error
checkError :: MonadBaseControl IO m => Maybe Int -> m () -> m ()
checkError code m = E.catch m handler
  where handler e@(CouchError c _) = unless (c == code) $ E.throwIO e

-- | Expect a couch error with the given code
assertRecvError :: (MonadIO m, MonadBaseControl IO m) => Maybe Int -> m a -> m ()
assertRecvError code v = checkError code $ v >> liftIO (assertFailure "was expecting a couch error")

-- | Check that an object in the database matches the given value.
checkLoad :: String -> A.Object -> CouchT IO ()
checkLoad n obj = do
    obj' <- couchGet n []
    lift $ assertBool "returned object does not match" $ isSubmapOf obj obj'

checkRevision :: String -> Revision -> CouchT IO ()
checkRevision n r = do
    r' <- couchRev n
    lift $ assertBool "returned revision does not match" $ r == r'

-- | Delete the given object, useful for the start of a test
clearObject :: String -> CouchT IO ()
clearObject n = checkError (Just 404) go
  where go = do rev <- couchRev n
                couchDelete n rev

newtype ArbitraryObject = ArbitraryObject { unArbObject :: A.Object }
    deriving (Show,Eq,A.FromJSON,A.ToJSON)

instance Arbitrary T.Text where
    arbitrary = liftM T.pack $ listOf $ elements $ ['a'..'z'] ++ ['A'..'Z'] ++ " 1234567890!@#$%^&*()+|"
    shrink "" = []
    shrink x  = [T.tail x]

arbBaseValue :: Gen A.Value
arbBaseValue = oneof [ A.String <$> arbitrary
                     , A.toJSON <$> (arbitrary :: Gen Integer)
                     , A.Bool <$> arbitrary
                     , return A.Null
                     ]

arbObject :: Bool -> Gen A.Object
arbObject onlyBase = do nkeys <- choose (3,15)
                        keys <- vectorOf nkeys arbitrary
                        vals <- vectorOf nkeys $ if onlyBase
                                                    then arbBaseValue
                                                    else frequency [ (8, arbBaseValue)
                                                                   , (1, A.Object <$> arbObject False)
                                                                   , (1, A.Array <$> arbArrayOfObj)
                                                                   ]

                        return $ M.fromList $ zip keys vals 

arbArrayOfObj :: Gen A.Array
arbArrayOfObj = do len <- choose (1,20)
                   vals <- vectorOf len (A.Object <$> arbObject False)
                   return $ V.fromList vals

instance Arbitrary ArbitraryObject where
    arbitrary = ArbitraryObject <$> arbObject True