{-# LANGUAGE FlexibleContexts #-}

module Thermoprint.Server.Fork
       ( ThreadManager
       , fork
       , cleanup
       , threadManager
       ) where

import Control.Monad.Reader.Class
import Control.Monad.Trans.Class
import Control.Monad.Trans.Reader (ReaderT, runReaderT)
import Control.Monad.Catch

import Control.Monad.IO.Class

import Control.Monad
import Control.Applicative
import Data.Maybe

import Data.Foldable

import Data.Map (Map)
import qualified Data.Map as Map

import Control.Concurrent
import Control.Concurrent.STM
import Control.Concurrent.STM.TVar (TVar)
import qualified Control.Concurrent.STM.TVar as T
import Control.Concurrent.STM.TSem (TSem)
import qualified Control.Concurrent.STM.TSem as S

data ThreadManager m = ThreadManager
  { fork    :: m () -> m ThreadId
  , cleanup :: m ()
  }

threadManager :: (MonadIO m, MonadMask m) => (m () -> m ThreadId) -> m (ThreadManager m)
threadManager f = do
  tVar <- newTVar Map.empty
  return ThreadManager
    { fork = \act -> do
        let
          unregisterSelf :: MonadIO m => m ()
          unregisterSelf = do
            tMap <- readTVar tVar
            tId  <- liftIO $ myThreadId
            modifyTVar' tVar $ Map.delete tId
            maybeM signalTSem $ Map.lookup tId tMap

        mask $ \unmask -> do
          tId <- f (unmask act `finally` unregisterSelf)
          modifyTVar' tVar =<< (Map.insert tId <$> newTSem 0)
          return tId
    , cleanup = liftIO $
        mapM_ (\(tId, s) -> killThread tId >> waitTSem s) . Map.toList =<< readTVar tVar
    }
  where
    atomically' :: MonadIO m => STM a -> m a
    atomically' = liftIO . atomically
  
    newTSem :: MonadIO m => Int -> m TSem
    newTSem = atomically' . S.newTSem
  
    waitTSem :: MonadIO m => TSem -> m ()
    waitTSem = atomically' . S.waitTSem
    
    signalTSem :: MonadIO m => TSem -> m ()
    signalTSem = atomically' . S.signalTSem

    newTVar :: MonadIO m => a -> m (TVar a)
    newTVar = atomically' . T.newTVar

    readTVar :: MonadIO m => TVar a -> m a
    readTVar = atomically' . T.readTVar

    modifyTVar' :: MonadIO m => TVar a -> (a -> a) -> m ()
    modifyTVar' t = atomically' . T.modifyTVar t

maybeM :: Applicative m => (a -> m ()) -> Maybe a -> m ()
maybeM = maybe $ pure ()