{-# 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 System.Mem.Weak
import qualified Data.Map as Map
import qualified Data.IntMap as IntMap
import Control.Exception.Lifted
import Control.Monad.Base
-- import Control.Monad.IO.Class
import qualified GHC.Conc as GHC
import Data.Time()
import Data.Time.Clock

data PerThread = PerThread
 { -- wkid :: Weak ThreadId
   lbl :: String
 , startTime :: UTCTime
 }
 deriving (Eq,Ord,Show) -- ,Data,Generic)

data GlobalState = GlobalState
 { threads :: Map.Map ThreadId PerThread -- IntMap.IntMap PerThread
 -- , uniqSource :: Int
 }

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


forkIO :: IO () -> IO ThreadId
forkIO = fork

fork :: MonadBaseControl IO m => m () -> m ThreadId
fork action = do
    t <- Raw.fork $ do
        -- wkid <- myThreadId >>= liftBase . mkWeakThreadId
        -- tid <- newUniq
        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
    putStrLn $ "labelThread "++s++" "++show tid
    modifyThreads $ Map.adjust (\pt -> pt { lbl = s }) tid

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

-- newUniq :: MonadBaseControl IO m => m Int
-- newUniq = do
--     modifyMVar globals (\st -> return (st { uniqSource = succ (uniqSource st) }, uniqSource st))
 
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) }
    putMVar globals (f' g)