module Lib where

import Control.Monad
import Control.Monad.STM (atomically)
import Control.Concurrent.Chan
import Control.Concurrent (threadDelay, newEmptyMVar, tryPutMVar, readMVar, isEmptyMVar, takeMVar)
import Ki qualified

{-
>>> thread "A" 1
A__1
>>> thread "A" 2
A__2
-}
thread :: String -> Int -> IO ()
thread s n = putStrLn (s ++ "__" ++ show n)

{-
>>> threadSequence "A" 3
A__1
A__2
A__3
3
-}
threadSequence :: String -> Int -> IO Int
threadSequence s n = do
                      forM [1..n] $ \x -> do
                        thread s x
                        threadDelay 1_000_000
                      pure n

{-
>>> concurrently (threadSequence "A" 5) (threadSequence "B" 7)
A__1
B__1
A__2
B__2
A__3
B__3
A__4
B__4
A__5
B__5
B__6
B__7
(5,7)
-}
concurrently :: IO a -> IO b -> IO (a, b)
concurrently a1 a2 =
    Ki.scoped \scope -> do
      thread1 <- Ki.fork scope a1
      threadDelay 200_000 -- to avoid both mixing up at the same line of stdout
      thread2 <- Ki.fork scope a2
      result1 <- atomically (Ki.await thread1)
      result2 <- atomically (Ki.await thread2)
      pure (result1, result2)

{-
>>> concurrentlyList [threadSequence "A" 5, threadSequence "B" 7]
A__1
B__1
A__2
B__2
A__3
B__3
A__4
B__4
A__5
B__5
B__6
B__7
-}
concurrentlyList :: [IO a] -> IO ()
concurrentlyList xs =
    Ki.scoped \scope -> do
      forM xs $ \x -> do
        threadDelay 200_000 -- note: we must delay threads between forking for not colliding on same line of stdout
        Ki.fork scope x
      atomically (Ki.awaitAll scope)

wrapper chan action = do
  result <- action
  _      <- writeChan chan result
  pure ()

{-
Even though we are using `forM` which should return `IO a`, our `concurrentlyList` doesn't return `IO [n]`
Therefore, we need to store the results in an `MVar`. A jar that is accessible by multiple threads.
-}
{-
>>> concurrentlyList' [threadSequence "A" 5, threadSequence "B" 7]
A__1
B__1
A__2
B__2
A__3
B__3
A__4
B__4
A__5
B__5
B__6
B__7
[5,7*** Exception: thread blocked indefinitely in an MVar operation
-}
concurrentlyList' xs = do
    resultChan <- newChan
    _ <- Ki.scoped \scope -> do
        forM xs $ \x -> do
          threadDelay 200_000 -- note: we must delay threads between forking for not colliding on same line of stdout
          Ki.fork scope $ wrapper resultChan x
        atomically (Ki.awaitAll scope)
    getChanContents resultChan