{-# LANGUAGE FlexibleContexts #-}
module Control.Concurrent.Lifted.Instrument
    ( module Control.Concurrent.Lifted
    , forkIO
    , fork
    , labelThread
    , threadsInformation
    , PerThread(..)
    ) where

import qualified Control.Concurrent.Lifted as Raw
import Control.Concurrent.Lifted hiding (fork)
import Control.Monad.Trans.Control
import System.IO.Unsafe
import qualified Data.Map.Strict  as Map
import Control.Exception.Lifted
import Control.Monad.Base
import qualified GHC.Conc as GHC
import Data.Time()
import Data.Time.Clock

data PerThread = PerThread
 { lbl :: String
 , startTime :: UTCTime
 }
 deriving (Eq,Ord,Show)

data GlobalState = GlobalState
 { threads :: Map.Map ThreadId PerThread
 }

globals :: MVar GlobalState
globals = unsafePerformIO $ newMVar $ GlobalState
    { threads = Map.empty
    }
{-# NOINLINE globals #-}


forkIO :: IO () -> IO ThreadId
forkIO = fork
{-# INLINE forkIO #-}

fork :: MonadBaseControl IO m => m () -> m ThreadId
fork action = do
    t <- Raw.fork $ do
        tid <- myThreadId
        tm <- liftBase getCurrentTime
        bracket_ (modifyThreads $ Map.insert tid (PerThread "" tm))
                 (modifyThreads $ Map.delete tid)
                 action
    return t

labelThread :: ThreadId -> String -> IO ()
labelThread tid s = do
    GHC.labelThread tid s
    modifyThreads $ Map.adjust (\pt -> pt { lbl = s }) tid
{-# INLINE labelThread #-}

threadsInformation :: IO [PerThread]
threadsInformation = do
    m <- threads <$> readMVar globals
    return $ Map.elems m


modifyThreads :: MonadBase IO m => (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
modifyThreads f = do
    g <- takeMVar globals
    let f' st = st { threads = f (threads st) }
        r = f' g
    putMVar globals r