module Main (main) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Control.Monad (forM)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Foldable (traverse_)
import Data.Functor ((<&>))
import Data.IORef
import Multitasking
import Test.Tasty
import Test.Tasty.HUnit

main :: IO ()
main =
  defaultMain $
    testGroup
      "tests"
      [ testCase "launch & wait" $ do
          multitask $ \coordinator -> do
            t1 <- start coordinator $ pure (4 :: Int)
            t2 <- start coordinator $ pure (4 :: Int)
            val1 <- await t1
            val2 <- await t2
            liftIO $ val1 @?= val2,
        testCase "launch many tasks" $ multitask $ \coordinator -> do
          counter <- newCounter 0
          tasks <- forM [1 .. 10000] $ \(_i :: Int) -> do
            start coordinator $ incrementCounter counter
          traverse_ await tasks
          value <- getCounter counter
          liftIO $ value @?= 10000,
        testCase "launch really many tasks" $ multitask $ \coordinator -> do
          counter <- newCounter 0
          tasks <- forM [1 .. 100000] $ \(_i :: Int) -> do
            start coordinator $ incrementCounter counter
          traverse_ await tasks
          value <- getCounter counter
          liftIO $ value @?= 100000,
        testCase "wait all" $ do
          ref <- newIORef (0 :: Int)
          multitask $ \coordinator -> do
            _ <- start coordinator $ atomicModifyIORef' ref (\a -> (a + 1, ()))
            _ <- start coordinator $ atomicModifyIORef' ref (\a -> (a + 1, ()))
            _ <- start coordinator $ atomicModifyIORef' ref (\a -> (a + 1, ()))
            awaitAll coordinator
            pure ()
          value <- readIORef ref
          value @?= 3,
        testCase "cancel task" $ do
          ref <- newIORef (0 :: Int)
          gate <- newGate
          multitask $ \coordinator -> do
            t1 <- start coordinator $ atomicModifyIORef' ref (\a -> (a + 1, ()))
            _ <- start coordinator $ awaitGate gate >> atomicModifyIORef' ref (\a -> (a + 1, ()))
            t3 <- start coordinator $ atomicModifyIORef' ref (\a -> (a + 1, ()))
            await t1
            await t3
            pure ()
          openGate gate
          value <- readIORef ref
          value @?= 2,
        testCase "cancel all tasks" $ do
          ref <- newIORef (0 :: Int)
          gate <- newGate
          multitask $ \coordinator -> do
            _ <- start coordinator $ awaitGate gate >> atomicModifyIORef' ref (\a -> (a + 1, ()))
            _ <- start coordinator $ awaitGate gate >> atomicModifyIORef' ref (\a -> (a + 1, ()))
            _ <- start coordinator $ awaitGate gate >> atomicModifyIORef' ref (\a -> (a + 1, ()))
            pure ()
          openGate gate
          value <- readIORef ref
          value @?= 0,
        testCase "wait canceled task 1" $ do
          _ <- assertKilledThread $ do
            t1 <- multitask $ \coordinator -> start coordinator waitForever
            await t1

          pure (),
        testCase "task error propagates to scope" $ do
          assertMultitaskException $ multitask $ \coordinator -> do
            gate <- newGate
            _ <- start coordinator $ undefined >> openGate gate
            _ <- start coordinator $ (pure ())
            awaitGate gate
            pure (),
        testCase "scope exception kills task" $ do
          killedRef <- newIORef False
          assertSomeException $ multitask $ \coordinator -> do
            gate <- newGate
            _ <-
              start coordinator $
                catch @SomeAsyncException
                  (openGate gate >> waitForever)
                  (\e -> writeIORef killedRef True >> throwIO e)
            awaitGate gate
            fail "I die"
          value <- readIORef killedRef
          value @?= True,
        testCase "task exception kills other task" $ do
          killedRef <- newIORef False
          assertMultitaskException $ multitask $ \coordinator -> do
            gate <- newGate
            _ <-
              start coordinator $
                catch @SomeAsyncException
                  (openGate gate >> waitForever)
                  (\e -> writeIORef killedRef True >> throwIO e)
            t <- start coordinator $ awaitGate gate >> fail "I die"
            () <- await t
            pure ()
          value <- readIORef killedRef
          value @?= True,
        testCase "race 2 actions" $ multitask $ \coordinator -> do
          gate1 <- newGate
          gate2 <- newGate
          task :: Task Int <- start coordinator $ raceTwo (awaitGate gate1 >> pure 1) (awaitGate gate2 >> pure 2)
          openGate gate2
          value <- await task
          liftIO $ value @?= 2,
        testCase "race many actions" $ multitask $ \coordinator -> do
          var <- newVariable (-1)
          task <-
            start coordinator $
              raceMany $
                [0 .. 10000] <&> \(i :: Int) -> do
                  awaitCondition (newCondition (== i) (readVariable var)) >> pure i
          writeVariable var 5000
          value <- await task
          liftIO $ value @?= 5000,
        testCase "race 2 tasks" $ multitask $ \coordinator -> do
          gate1 <- newGate
          gate2 <- newGate
          t1 <- start coordinator $ awaitGate gate1 >> pure 1
          t2 <- start coordinator $ awaitGate gate2 >> pure 2
          t3 :: Task Int <- start coordinator $ raceTwo (await t1) (await t2)
          openGate gate2
          value <- await t3
          liftIO $ value @?= 2,
        testCase "race with finished task" $ multitask $ \coordinator -> do
          gate1 <- newGate
          gate2 <- newGate
          t1 <- start coordinator $ awaitGate gate1 >> pure (1 :: Int)
          t2 <- start coordinator $ awaitGate gate2 >> pure (2 :: Int)
          openGate gate2
          _ <- await t2
          t3 <- start coordinator $ raceTwo (await t1) (await t2)
          value <- await t3
          liftIO $ value @?= 2,
        testCase "raceMaybe 2 tasks 1" $ do
          result <- raceTwoMaybe (pure Nothing) (threadDelay 100 >> pure (Just ()))
          liftIO $ result @?= Just (),
        testCase "raceMaybe 2 tasks 2" $ do
          result <- raceTwoMaybe (pure (Just ())) (threadDelay 100 >> pure Nothing)
          liftIO $ result @?= Just (),
        testCase "raceMaybe 2 tasks 3" $ do
          result :: Maybe Int <- raceTwoMaybe (pure Nothing) (threadDelay 100 >> pure Nothing)
          liftIO $ result @?= Nothing,
        testCase "raceMaybe many tasks" $ do
          result :: Maybe Int <- raceManyMaybe [pure Nothing, threadDelay 100 >> pure Nothing]
          liftIO $ result @?= Nothing,
        testCase "raceMaybe many tasks" $ do
          result <- raceManyMaybe [pure Nothing, threadDelay 100 >> pure (Just ())]
          liftIO $ result @?= Just (),
        testCase "raceMaybe many tasks" $ do
          result <- raceManyMaybe [pure (Just ()), threadDelay 10000 >> pure undefined]
          liftIO $ result @?= Just (),
        testCase "timeout" $ do
          result :: Maybe Int <- timeout (fromMilliseconds 10) waitForever
          liftIO $ result @?= Nothing,
        testCase "max concurrency" $ multitask $ \coordinator -> do
          limit <- maxConcurrentTasks 3
          counter <- newCounter 0
          forM_ [0 .. 100] $ \i ->
            start coordinator $
              throttle limit $
                incrementCounter counter *> waitDuration (fromMilliseconds 10)
          waitDuration (fromMilliseconds 15)
          result <- getCounter counter
          liftIO $ result @?= 6,
        testCase "token bucket" $ multitask $ \coordinator -> do
          limit <- tokenBucket coordinator (fromMilliseconds 3) 3
          counter <- newCounter 0
          forM_ [0 .. 100] $ \i ->
            start coordinator $
              throttle limit $
                incrementCounter counter
          waitDuration (fromMilliseconds 13)
          result <- getCounter counter
          liftIO $ result @?= 7
      ]

assertSomeException :: IO a -> IO ()
assertSomeException action = do
  result <- try @SomeException action
  case result of
    Left _ -> return ()
    Right _ -> assertFailure "Expected SomeException, but none thrown"

assertKilledThread :: IO a -> IO ()
assertKilledThread action = do
  result <- try @SomeAsyncException action
  case result of
    Left _ -> return ()
    Right _ -> assertFailure "Expected a killed thread, but not killed"

assertMultitaskException :: IO a -> IO ()
assertMultitaskException action = do
  result <- try @SomeException action
  case result of
    Left _ -> return ()
    Right _ -> assertFailure "Expected an exception, but none thrown"