{-#LANGUAGE DeriveDataTypeable#-}
{-#LANGUAGE BangPatterns#-}

module UnliftIO.PooledAsyncSpec (spec) where

import Test.Hspec
import Control.Concurrent
import Data.List (sort)
import Test.QuickCheck
import qualified Data.Set as Set
import Data.Functor ((<$>))
import UnliftIO

data MyPooledException = PoolHellException
                         deriving (Show, Typeable)

instance Exception MyPooledException

-- | Strip out duplicates. (Taken from rio)
nubOrd :: Ord a => [a] -> [a]
nubOrd =
  loop Set.empty
  where
    loop _ [] = []
    loop !s (a:as)
      | a `Set.member` s = loop s as
      | otherwise = a : loop (Set.insert a s) as

spec :: Spec
spec = do
  let exAction :: Int -> IO Int
      exAction x = do
        if (x == 2) then throwIO PoolHellException else return ()
        return x

      action :: Int -> IO ThreadId
      action x = do
         threadDelay (2 * 10^5)
         myThreadId

      myVar :: IO (TVar Int)
      myVar = atomically $ newTVar 0

      maxTVar :: Int -> TVar Int -> IO ()
      maxTVar cval tvar = do
         atomically $ do
           v <- readTVar tvar
           if cval >= v
           then writeTVar tvar cval
           else return ()

      poolException :: Selector MyPooledException
      poolException = const True

  describe "pooled mapConcurrencyN" $ do
    it "Throws exception properly" $ do
       (pooledMapConcurrentlyN 5 exAction [1..5]) `shouldThrow` poolException

    it "total thread should be >= 1" $ do
       (pooledMapConcurrentlyN 0 action [1..5]) `shouldThrow` anyErrorCall

    it "should not spawn more than five threads for five concurrent tasks" $ do
       xs <- (pooledMapConcurrentlyN 5 action [1..5])
       (length $ nubOrd xs) `shouldSatisfy` (<= (5 :: Int))

    it "should not spawn more than three threads for five concurrent tasks" $ do
       xs <- (pooledMapConcurrentlyN 3 action [1..5])
       (length $ nubOrd xs) `shouldSatisfy` (<= (3 :: Int))

    it "should spawn only one thread" $ do
       xs <- (pooledMapConcurrentlyN 1 action [1..5])
       (length $ nubOrd xs) `shouldBe` 1

    it "never uses more than the given number of pools and doesn't miss any return values" $
        forAllShrink ((+ 1) . abs <$> arbitrary) (filter (>= 1) . shrink) $ \threads ->
            property $ \list -> do
                threadIdsVar <- newTVarIO []
                let go :: Int -> IO Int
                    go i = do
                        tid <- myThreadId
                        atomically $ modifyTVar threadIdsVar (tid :)
                        return i
                list' <- pooledMapConcurrentlyN threads go list
                sort list' `shouldBe` sort list
                tids <- readTVarIO threadIdsVar
                length (nubOrd tids) `shouldSatisfy` (<= threads)

  describe "pooled mapConcurrencyN_" $ do
    it "Throws exception properly" $ do
       (pooledMapConcurrentlyN_ 5 exAction [1..5]) `shouldThrow` poolException

    it "total thread should be >= 1" $ do
       (pooledMapConcurrentlyN_ 0 action [1..5]) `shouldThrow` anyErrorCall

    it "find proper maximum value" $ do
       var <- myVar
       xs <- (pooledMapConcurrentlyN_ 5 (\x -> maxTVar x var) [1..5])
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 5

    it "find proper maximum value with 2 threads" $ do
       var <- myVar
       xs <- (pooledMapConcurrentlyN_ 2 (\x -> maxTVar x var) [1..5])
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 5

    it "find proper maximum value with 1 threads" $ do
       var <- myVar
       xs <- (pooledMapConcurrentlyN_ 1 (\x -> maxTVar x var) [1..5])
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 5

    it "make sure activity is happening in different threads" $ do
       let myThreads :: IO (TVar [ThreadId])
           myThreads = atomically $ newTVar []

           collectThreads :: TVar [ThreadId] -> IO ()
           collectThreads threadVar = do
             tid <- myThreadId
             atomically $ do
               tvar <- readTVar threadVar
               writeTVar threadVar (tid:tvar)
             threadDelay $ 2 * 10^5

       tid <- myThreads
       xs <- pooledMapConcurrentlyN_ 5 (\_ -> collectThreads tid) [1..5]
       tids <- atomically $ readTVar tid
       (length $ nubOrd tids) `shouldSatisfy` (<= 5)

    it "Not more than 5 threads will be spawned even if pooling is set to 8 " $ do
       let myThreads :: IO (TVar [ThreadId])
           myThreads = atomically $ newTVar []

           collectThreads :: TVar [ThreadId] -> IO ()
           collectThreads threadVar = do
             tid <- myThreadId
             atomically $ do
               tvar <- readTVar threadVar
               writeTVar threadVar (tid:tvar)
             threadDelay $ 2 * 10^5

       tid <- myThreads
       xs <- pooledMapConcurrentlyN_ 8 (\_ -> collectThreads tid) [1..5]
       tids <- atomically $ readTVar tid
       (length $ nubOrd tids) `shouldSatisfy` (<= 5)

  describe "replicate concurrencyN" $ do
    it "Throws exception properly" $ do
       (pooledReplicateConcurrentlyN 5 1 (exAction 2)) `shouldThrow` poolException

    it "total thread should be >= 1" $ do
       (pooledReplicateConcurrentlyN 0 1 (action 1)) `shouldThrow` anyErrorCall

    it "Read tvar value should be 100" $ do
       var <- myVar
       xs <- (pooledReplicateConcurrentlyN 5 5 (maxTVar 100 var))
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 100

    it "should not spawn more than five threads for five concurrent tasks" $ do
       xs <- (pooledReplicateConcurrentlyN 5 5 (action 1))
       (length $ nubOrd xs) `shouldSatisfy` (<= (5 :: Int))

    it "should not spawn more than three threads for five concurrent tasks" $ do
       xs <- (pooledReplicateConcurrentlyN 3 5 (action 1))
       (length $ nubOrd xs) `shouldSatisfy` (<= (3 :: Int))

    it "should spawn only one thread" $ do
       xs <- (pooledReplicateConcurrentlyN 1 5 (action 1))
       (length $ nubOrd xs) `shouldBe` 1

    it "should give empty list" $ do
       xs <- (pooledReplicateConcurrentlyN 3 0 (action 1))
       xs `shouldBe` []

    it "should give empty list for -ve count" $ do
       xs <- (pooledReplicateConcurrentlyN 3 (-3) (action 1))
       xs `shouldBe` []

  describe "pooled replicateConcurrencyN_" $ do
    it "Throws exception properly" $ do
       (pooledReplicateConcurrentlyN_ 5 1 (exAction 2)) `shouldThrow` poolException

    it "total thread should be >= 1" $ do
       (pooledReplicateConcurrentlyN_ 0 2 (action 1)) `shouldThrow` anyErrorCall

    it "find proper maximum value" $ do
       var <- myVar
       pooledReplicateConcurrentlyN_ 5 3 (maxTVar 200 var)
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 200

    it "Should be initial value" $ do
       var <- myVar
       pooledReplicateConcurrentlyN_ 5 (-2) (maxTVar 200 var)
       newVar <- atomically $ readTVar var
       atomically $ writeTVar var 0
       newVar `shouldBe` 0