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

import qualified Control.Concurrent.Lifted as Raw
import Control.Concurrent.Lifted hiding (fork,forkOS)
import Control.Exception (fromException)
import Control.Monad.Trans.Control
import GHC.Stack
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
import DPut
import DebugTag


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

data GlobalState = GlobalState
 { threads :: !(Map.Map ThreadId PerThread)
 , reportException :: String -> IO ()
 }

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

forkLabeled :: HasCallStack => String -> IO () -> IO ThreadId
forkLabeled lbl action = do
    t <- instrumented GHC.forkIO action
    labelThread t lbl
    return t
{-# INLINE forkLabeled #-}

forkOSLabeled :: HasCallStack => String -> IO () -> IO ThreadId
forkOSLabeled lbl action = do
    t <- instrumented Raw.forkOS action
    labelThread t lbl
    return t
{-# INLINE forkOSLabeled #-}

forkIO :: HasCallStack => IO () -> IO ThreadId
forkIO = instrumented GHC.forkIO
{-# INLINE forkIO #-}

forkOS :: ( HasCallStack, MonadBaseControl IO m ) => m () -> m ThreadId
forkOS = instrumented Raw.forkOS
{-# INLINE forkOS #-}

fork :: ( HasCallStack, MonadBaseControl IO m ) => m () -> m ThreadId
fork = instrumented Raw.fork
{-# INLINE fork #-}

shortCallStack :: [([Char], SrcLoc)] -> String
shortCallStack []          = ""
shortCallStack ((_,loc):_) = (srcLocFile loc) ++ ":" ++ show (srcLocStartLine loc)

defaultLabel :: CallStack -> String
defaultLabel stack = case getCallStack stack of
    _ : sites  -> shortCallStack sites
    sites -> shortCallStack sites

instrumented :: ( HasCallStack, MonadBaseControl IO m ) =>
                (m () -> m ThreadId) -> m () -> m ThreadId
instrumented rawFork action = do
    mvar <- newEmptyMVar
    t <- rawFork $ do
        tid <- myThreadId
        tm <- liftBase getCurrentTime
        bracket_ (modifyThreads $! \ts -> Map.union ts (Map.singleton tid (PerThread (defaultLabel callStack) tm)))
                 (return ())
            $ do catch action $ \e -> case fromException e of
                    Just ThreadKilled -> return ()
                    Nothing           -> liftBase $ do
                        bracket (takeMVar globals)
                                (\g -> do
                                    let l = concat [ show e
                                                   , " ("
                                                   , maybe "" lbl $ Map.lookup tid (threads g)
                                                   , ")"
                                                   ]
                                    foldr seq (return ()) l
                                    putMVar globals $! g { threads = Map.insert tid (PerThread l tm) $ threads g }
                                    throwIO e)
                                (\g -> do
                                    let l = concat [ show e
                                                   , " ("
                                                   , maybe "" lbl $ Map.lookup tid (threads g)
                                                   , ")"
                                                   ]
                                    reportException g l)
                 -- Remove the thread only if it terminated normally or was killed.
                 takeMVar mvar
                 modifyThreads $! Map.delete tid
    liftBase $ labelThread_ t (defaultLabel callStack)
    putMVar mvar ()
    return t

labelThread_ :: ThreadId -> String -> IO ()
labelThread_ tid s = do
    foldr seq (return ()) s
    GHC.labelThread tid s
    tm <- liftBase getCurrentTime
    let updateIt (Just pt) = Just $ pt { lbl = s }
        updateIt Nothing   = Just $ PerThread s tm
    modifyThreads $! Map.alter updateIt tid

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

threadsInformation :: IO [(ThreadId,PerThread)]
threadsInformation = do
    m <- threads <$> readMVar globals
    return $ Map.toList m


modifyThreads :: MonadBaseControl IO m =>
                    (Map.Map ThreadId PerThread -> Map.Map ThreadId PerThread) -> m ()
modifyThreads f = do
    let f' st = st { threads = f (threads st) }
    bracket (takeMVar globals)
            (\g -> putMVar globals $! f' g)
            (\g -> return ())