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

module Main where

import           Control.Applicative
import qualified Control.Exception.Lifted as E
import           Control.Monad
import           Control.Monad.IO.Class (MonadIO, liftIO)
import           Control.Monad.Trans.Class (lift)
import           Control.Monad.Trans.Control (MonadBaseControl)
import           Control.Monad.Trans.Reader
import           Data.Aeson ((.=))
import qualified Data.Aeson as A
import qualified Data.ByteString.UTF8 as BU8
import           Data.Enumerator hiding (map, length, head, run, drop)
import qualified Data.Enumerator.List as EL
import           Data.List (find, deleteBy, nubBy)
import qualified Data.HashMap.Lazy as M
import           Data.Maybe (fromJust)
import qualified Data.Text as T
import qualified Data.Vector as V

import Database.CouchDB.Enumerator

import Test.Framework (defaultMain, 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)

main :: IO ()
main = defaultMain tests

tests :: [Test]
tests = [ testCouchCase     "basic connection"      connectTest
        , testCouchCase     "basic error"           missingObjectTest
        , testCouchProperty "single insert" (1,7)   insertTest
        , testCouchCase     "basic delete"          deleteTest
        , testCouchProperty "view"  (12,13)         views
        ]
   
----------------------------------------------------------------------------------
--- Test Helpers
----------------------------------------------------------------------------------

type CouchT m a = ReaderT CouchConnection m a

testCouch :: CouchT IO a -> IO ()
testCouch c = withCouchConnection "172.16.5.2" 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'

-- | Delete the given object, useful for the start of a test
clearObject :: String -> CouchT IO ()
clearObject n = checkError (Just 404) go
  where go = do obj <- couchGet n []
                unless (member "_rev" obj) $ fail "_rev is missing"
                let (A.String rev) = fromJust $ M.lookup "_rev" obj
                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

----------------------------------------------------------------------------------
--- Base Tests
----------------------------------------------------------------------------------

connectTest :: CouchT IO ()
connectTest = do
    v <- couchGet "" []
    lift $ assertObjMember "db_name" (assertStr "testcouchenum") v

missingObjectTest :: CouchT IO ()
missingObjectTest = assertRecvError (Just 404) $ couchGet "jaosihaweoghaweiouhawef" []

insertTest :: [(Int,ArbitraryObject,ArbitraryObject)] -> CouchT IO ()
insertTest objs = do
    let objs' = nubBy (\(a,_,_) (b,_,_) -> a == b) objs
    let keys  = map (("otest"++) . show . (\(a,_,_) -> a)) objs'
    let vals1 = map (\(_,ArbitraryObject a,_) -> a) objs'
    let vals2 = map (\(_,_,ArbitraryObject a) -> a) objs'

    mapM_ clearObject keys

    rev <- forM (zip keys vals1) $ \(k,o) -> couchPut k [] o
    forM_ (zip3 rev keys vals1) $ \(r,k,o) ->
        checkLoad k $ M.insert "_rev" (A.toJSON r) o

    rev2 <- forM (zip3 rev keys vals2) $ \(r,k,o) ->
        couchPut k [] $ M.insert "_rev" (A.toJSON r) o

    forM_ (zip3 rev2 keys vals2) $ \(r,k,o) ->
        checkLoad k $ M.insert "_rev" (A.toJSON r) o

deleteTest :: CouchT IO ()
deleteTest = do
    (ArbitraryObject obj) <- liftM (head . drop 5) $ lift $ sample' arbitrary

    clearObject "deltest"

    rev <- couchPut "deltest" [] obj
    checkLoad "deltest" obj
    couchDelete "deltest" rev
    assertRecvError (Just 404) $ couchGet "deltest" []

----------------------------------------------------------------------------------
--- View Tests
----------------------------------------------------------------------------------

checkEqual :: (Monad m) => [A.Object] -> [A.Object] -> Iteratee a m ()
checkEqual []     []    = return ()
checkEqual []     (x:_) = error $ "Extra object in list 2  " ++ show (A.encode x)
checkEqual (x:xs) lst2  = case find (isSubmapOf x) lst2 of
                            Nothing -> error $ "Unable to find " ++ show (A.encode $ A.Object x)
                            Just _  -> checkEqual xs $ deleteBy isSubmapOf x lst2

assertViewRet :: (MonadIO m) => [A.Object] -> Enumerator A.Object m () -> m ()
assertViewRet lst e = run_ (e $$ EL.consume >>= checkEqual lst)

addKeys :: Int -> Int -> Int -> ArbitraryObject -> A.Object
addKeys u g t (ArbitraryObject o) = o `M.union` M.fromList [ ("user", A.toJSON u)
                                                           , ("group", A.toJSON g)
                                                           , ("otype", A.toJSON t)
                                                           ]

addView :: CouchT IO ()
addView = couchPut_ "_design/dataviews" [] viewObj where
   viewObj = A.object 
        [ "language"    .= ("javascript" :: T.Text)
        , "views"       .= A.object
            [ "bytype"  .= A.object
                [ "map" .= ("function(doc) {\
                              \   emit([doc.user,doc.group,doc.otype], doc); \
                              \}" :: T.Text)
                ]
            ]
        ]

queryByType :: Integer -> Integer -> Integer -> Enumerator A.Object (ReaderT CouchConnection IO) b
queryByType u g t = couchView path query $= extractViewValue
    where path  = "dataviews/_view/bytype"
          key   = "[" ++ show u ++ "," ++ show g ++ "," ++ show t ++ "]"
          query = [(BU8.fromString "key", Just $ BU8.fromString key)]

queryByGroup :: Integer -> Integer -> Enumerator A.Object (ReaderT CouchConnection IO) b
queryByGroup u g = couchView path query $= extractViewValue
    where path  = "dataviews/_view/bytype"
          skey  = "[" ++ show u ++ "," ++ show g ++ "]"
          ekey  = "[" ++ show u ++ "," ++ show g ++ ",{}]"
          query = [ (BU8.fromString "startkey", Just $ BU8.fromString skey)
                  , (BU8.fromString "endkey"  , Just $ BU8.fromString ekey)
                  ]

views :: [ArbitraryObject] -> CouchT IO ()
views lst = do
    let (group1,x1) = splitAt 3 lst
        (group2,x2) = splitAt 3 x1
        (group3,x3) = splitAt 3 x2
        (group4,_ ) = splitAt 3 x3

        g1key = map (("view"++) . show) ([0..2] :: [Int])
        g1obj = map (addKeys 0 0 0) group1
        g2key = map (("view"++) . show) ([5..7] :: [Int])
        g2obj = map (addKeys 0 0 1) group2
        g3key = map (("view"++) . show) ([10..12] :: [Int])
        g3obj = map (addKeys 0 1 0) group3
        g4key = map (("view"++) . show) ([15..17] :: [Int])
        g4obj = map (addKeys 1 0 0) group4

    mapM_ clearObject $ g1key ++ g2key ++ g3key ++ g4key

    checkError (Just 409) addView

    forM_ (zip g1key g1obj ++ zip g2key g2obj ++ zip g3key g3obj ++ zip g4key g4obj) $ \(k,o) ->
        couchPut_ k [] o

    assertViewRet []    $ queryByType 0 0 2
    assertViewRet []    $ queryByType 0 2 0
    assertViewRet []    $ queryByType 2 0 0

    assertViewRet g1obj $ queryByType 0 0 0
    assertViewRet g2obj $ queryByType 0 0 1
    assertViewRet g3obj $ queryByType 0 1 0
    assertViewRet g4obj $ queryByType 1 0 0

    assertViewRet (g1obj ++ g2obj) $ queryByGroup 0 0
    assertViewRet g3obj $ queryByGroup 0 1
    assertViewRet g4obj $ queryByGroup 1 0

    assertViewRet [] $ queryByGroup 0 2
    assertViewRet [] $ queryByGroup 2 0

-- vim: set expandtab:set tabstop=4: