{-# LANGUAGE CPP #-}
{-# LANGUAGE TupleSections #-}

module Test.Sandwich.Interpreters.StartTree (
  startTree
  , runNodesSequentially
  , markAllChildrenWithResult
  ) where


import Control.Concurrent.MVar
import qualified Control.Exception as E
import Control.Monad
import Control.Monad.IO.Class
import Control.Monad.IO.Unlift
import Control.Monad.Logger
import Control.Monad.Trans.Reader
import qualified Data.ByteString.Char8 as BS8
import Data.IORef
import qualified Data.List as L
import Data.Sequence hiding ((:>))
import qualified Data.Set as S
import Data.String.Interpolate
import qualified Data.Text as T
import qualified Data.Text.IO as T
import Data.Time
import Data.Time.Clock.POSIX
import Data.Typeable
import GHC.IO (unsafeUnmask)
import GHC.Stack
import System.FilePath
import System.IO
import Test.Sandwich.Formatters.Print
import Test.Sandwich.Formatters.Print.CallStacks
import Test.Sandwich.Formatters.Print.FailureReason
import Test.Sandwich.Formatters.Print.Logs
import Test.Sandwich.Formatters.Print.Printing
import Test.Sandwich.Interpreters.RunTree.Logging
import Test.Sandwich.Interpreters.RunTree.Util
import Test.Sandwich.ManagedAsync
import Test.Sandwich.RunTree
import Test.Sandwich.TestTimer
import Test.Sandwich.Types.RunTree
import Test.Sandwich.Types.Spec
import Test.Sandwich.Types.TestTimer
import Test.Sandwich.Util
import UnliftIO.Async
import UnliftIO.Concurrent (threadDelay)
import UnliftIO.Directory
import UnliftIO.Exception
import UnliftIO.STM


baseContextFromCommon :: RunNodeCommonWithStatus s l t -> BaseContext -> BaseContext
baseContextFromCommon (RunNodeCommonWithStatus {..}) bc@(BaseContext {}) =
  bc { baseContextPath = runTreeFolder }

startTree :: (MonadIO m, HasBaseContext context) => RunNode context -> context -> m (Async Result)
startTree node@(RunNodeBefore {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    let label = runTreeLabel <> " (setup)"
    (timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label (runExampleM runNodeCommon label runNodeBefore ctx runTreeLogs (Just [i|Exception in before '#{runTreeLabel}' handler|]))) >>= \case
      (result@(Failure fr@(Pending {})), _setupStartTime, setupFinishTime) -> do
        markAllChildrenWithResult runNodeChildren ctx (Failure fr)
        return (result, mkSetupTimingInfo setupFinishTime)
      (result@(Failure fr), _setupStartTime, setupFinishTime) -> do
        markAllChildrenWithResult runNodeChildren ctx (Failure $ GetContextException Nothing (SomeExceptionWithEq $ toException fr))
        return (result, mkSetupTimingInfo setupFinishTime)
      (Success, _setupStartTime, setupFinishTime) -> do
        void $ runNodesSequentially runNodeChildren ctx
        return (Success, mkSetupTimingInfo setupFinishTime)
      (Cancelled, _setupStartTime, setupFinishTime) -> do
        return (Cancelled, mkSetupTimingInfo setupFinishTime)
      (DryRun, _setupStartTime, setupFinishTime) -> do
        void $ runNodesSequentially runNodeChildren ctx
        return (DryRun, mkSetupTimingInfo setupFinishTime)
startTree node@(RunNodeAfter {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    result <- liftIO $ newIORef (Success, emptyExtraTimingInfo)
    -- We use the 'finally' from @base@; see the comment below about 'bracket'
    E.finally (void $ runNodesSequentially runNodeChildren ctx)
              (do
                  let label = runTreeLabel <> " (teardown)"
                  (ret, teardownStartTime, _teardownFinishTime) <- timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label $
                    runExampleM runNodeCommon label runNodeAfter ctx runTreeLogs (Just [i|Exception in after '#{runTreeLabel}' handler|])
                  writeIORef result (ret, mkTeardownTimingInfo teardownStartTime)
              )
    liftIO $ readIORef result
startTree node@(RunNodeIntroduce {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    result <- liftIO $ newIORef (Success, emptyExtraTimingInfo)
    -- The bracket from @base@ uses 'mask', while the one from @unliftio@ uses 'uninterruptibleMask'.
    -- We want to err on the side of avoiding deadlocks so we use the @base@ version.
    -- This also plays nicer with the 'withMaybeWarnOnLongExecution' and 'withMaybeCancelOnLongExecution' functions.
    -- See the discussion at https://github.com/fpco/safe-exceptions/issues/3
    let opts = baseContextOptions (getBaseContext ctx)
    liftIO $ E.bracket
      (do
          let asyncExceptionResult e = Failure $ GotAsyncException Nothing (Just [i|introduce #{runTreeLabel} alloc handler got async exception|]) (SomeAsyncExceptionWithEq e)
          let label = runTreeLabel <> " (setup)"
          getCurrentTime >>= \t -> emitEvent opts t runTreeId runTreeLabel EventSetupStarted
          flip withException (\(e :: SomeAsyncException) -> markAllChildrenWithResult runNodeChildrenAugmented ctx (asyncExceptionResult e)) $ do
            ret <- timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label $
              runExampleM' runNodeCommon label runNodeAlloc ctx runTreeLogs (Just [i|Failure in introduce '#{runTreeLabel}' allocation handler|])
            getCurrentTime >>= \t -> emitEvent opts t runTreeId runTreeLabel EventSetupFinished
            return ret
      )
      (\(ret, setupStartTime, setupFinishTime) -> case ret of
          Left failureReason -> writeIORef result (Failure failureReason, mkSetupTimingInfo setupStartTime)
          Right intro -> do
            teardownStartTime <- getCurrentTime
            addTeardownStartTimeToStatus runTreeStatus teardownStartTime
            emitEvent opts teardownStartTime runTreeId runTreeLabel EventTeardownStarted
            let label = runTreeLabel <> " (teardown)"
            (ret', _, _) <- timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label $
              runExampleM runNodeCommon label (runNodeCleanup intro) ctx runTreeLogs (Just [i|Failure in introduce '#{runTreeLabel}' cleanup handler|])
            getCurrentTime >>= \t -> emitEvent opts t runTreeId runTreeLabel EventTeardownFinished
            writeIORef result (ret', ExtraTimingInfo (Just setupFinishTime) (Just teardownStartTime))
      )
      (\(ret, _setupStartTime, setupFinishTime) -> do
          addSetupFinishTimeToStatus runTreeStatus setupFinishTime
          case ret of
            Left failureReason@(Pending {}) -> do
              -- TODO: add note about failure in allocation
              markAllChildrenWithResult runNodeChildrenAugmented ctx (Failure failureReason)
            Left failureReason -> do
              -- TODO: add note about failure in allocation
              markAllChildrenWithResult runNodeChildrenAugmented ctx (Failure $ GetContextException Nothing (SomeExceptionWithEq $ toException failureReason))
            Right intro -> do
              -- Special hack to modify the test timer profile via an introduce, without needing to track it everywhere.
              -- It would be better to track the profile at the type level
              let ctxFinal = case cast intro of
                    Just (TestTimerProfile t) -> modifyBaseContext ctx (\bc -> bc { baseContextTestTimerProfile = t })
                    Nothing -> ctx

              void $ runNodesSequentially runNodeChildrenAugmented ((LabelValue intro) :> ctxFinal)
      )
    readIORef result
startTree node@(RunNodeIntroduceWith {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  didRunWrappedAction <- liftIO $ newIORef (Left (), emptyExtraTimingInfo)
  let label = runTreeLabel <> " (body)"
  runInAsync node ctx $ do
    let wrappedAction = do
          didAllocateVar <- liftIO $ newIORef False
          let handler e = do
                liftIO (readIORef didAllocateVar) >>= \case
                  False -> do
                    -- We didn't successfully allocate, so mark the children below this point as failed due to a
                    -- context issue.
                    markAllChildrenWithResult runNodeChildrenAugmented ctx $ case fromException e of
                      Just fr@(Pending {}) -> Failure fr
                      _ -> Failure $ GetContextException Nothing (SomeExceptionWithEq e)
                  True ->
                    -- Otherwise, let their existing statuses stand.
                    return ()

          flip withException handler $ do
            beginningCleanupVar <- liftIO $ newIORef Nothing

            -- Record a start event in the test timing, if configured
            -- TODO: how to do this? The speedscope viewer crashes if you don't properly nest events
            -- let tt = if runTreeRecordTime then getTestTimer (getBaseContext ctx) else NullTestTimer
            -- let setupLabel = runTreeLabel <> " (setup)"
            -- let teardownLabel = runTreeLabel <> " (teardown)"
            -- handleStartEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack setupLabel)

            let opts = baseContextOptions (getBaseContext ctx)
            liftIO getCurrentTime >>= \t -> liftIO $ emitEvent opts t runTreeId runTreeLabel EventSetupStarted
            results <- runNodeIntroduceAction $ \intro -> do
              -- Record the end event in the test timing
              -- TODO: do we need to deal with exceptions here and in teardown? I think we we fail to emit the
              -- end event, the time profile won't be viewable.
              -- handleEndEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack setupLabel)
              setupFinishTime <- liftIO getCurrentTime
              addSetupFinishTimeToStatus runTreeStatus setupFinishTime
              liftIO $ emitEvent opts setupFinishTime runTreeId runTreeLabel EventSetupFinished

              liftIO $ writeIORef didAllocateVar True

              (results, _, _) <- timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label $
                liftIO $ runNodesSequentially runNodeChildrenAugmented (LabelValue intro :> ctx)

              teardownStartTime <- liftIO getCurrentTime
              addTeardownStartTimeToStatus runTreeStatus teardownStartTime
              liftIO $ emitEvent opts teardownStartTime runTreeId runTreeLabel EventTeardownStarted

              liftIO $ writeIORef beginningCleanupVar (Just teardownStartTime)
              -- handleStartEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack teardownLabel)
              liftIO $ writeIORef didRunWrappedAction (Right results, mkSetupTimingInfo setupFinishTime)
              return results

            liftIO (readIORef beginningCleanupVar) >>= \case
              Nothing -> return ()
              Just teardownStartTime -> do
                -- handleEndEvent tt (baseContextTestTimerProfile (getBaseContext ctx)) (T.pack teardownLabel)
                liftIO getCurrentTime >>= \t -> liftIO $ emitEvent opts t runTreeId runTreeLabel EventTeardownFinished
                liftIO $ modifyIORef' didRunWrappedAction $ \(ret, timingInfo) ->
                  (ret, timingInfo { teardownStartTime = Just teardownStartTime })

            return results

          liftIO (readIORef didRunWrappedAction) >>= \case
            (Left (), timingInfo) -> return (Failure $ Reason Nothing [i|introduceWith '#{runTreeLabel}' handler didn't call action|], timingInfo)
            (Right _, timingInfo) -> return (Success, timingInfo)
    runExampleM' runNodeCommon label wrappedAction ctx runTreeLogs (Just [i|Exception in introduceWith '#{runTreeLabel}' handler|]) >>= \case
      Left err -> return (Failure err, emptyExtraTimingInfo)
      Right x -> pure x
startTree node@(RunNodeAround {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  didRunWrappedAction <- liftIO $ newIORef (Left (), emptyExtraTimingInfo)
  runInAsync node ctx $ do
    let opts = baseContextOptions (getBaseContext ctx)
    let wrappedAction = do
          didStartChildren <- liftIO $ newIORef False
          let handler e =
                -- If the handler failed before the children were started, mark them as failed
                -- due to a context issue. (Don't touch the node's own status here: that would
                -- race with the authoritative write in 'runInAsync'.) Otherwise let the
                -- children's existing statuses stand.
                liftIO (readIORef didStartChildren) >>= \case
                  False -> markAllChildrenWithResult runNodeChildren ctx $ case fromException e of
                    Just fr@(Pending {}) -> Failure fr
                    _ -> Failure $ GetContextException Nothing (SomeExceptionWithEq e)
                  True -> return ()

          flip withException handler $ do
            liftIO getCurrentTime >>= \t -> liftIO $ emitEvent opts t runTreeId runTreeLabel EventSetupStarted
            runNodeActionWith $ do
              setupFinishTime <- liftIO getCurrentTime
              addSetupFinishTimeToStatus runTreeStatus setupFinishTime
              liftIO $ emitEvent opts setupFinishTime runTreeId runTreeLabel EventSetupFinished
              liftIO $ writeIORef didStartChildren True
              results <- liftIO $ runNodesSequentially runNodeChildren ctx
              liftIO getCurrentTime >>= \t -> liftIO $ emitEvent opts t runTreeId runTreeLabel EventTeardownStarted
              liftIO $ writeIORef didRunWrappedAction (Right results, mkSetupTimingInfo setupFinishTime)
              return results

          liftIO getCurrentTime >>= \t -> liftIO $ emitEvent opts t runTreeId runTreeLabel EventTeardownFinished
          (liftIO $ readIORef didRunWrappedAction) >>= \case
            (Left (), timingInfo) -> return (Failure $ Reason Nothing [i|around '#{runTreeLabel}' handler didn't call action|], timingInfo)
            (Right _, timingInfo) -> return (Success, timingInfo)
    runExampleM' runNodeCommon runTreeLabel wrappedAction ctx runTreeLogs (Just [i|Exception in around '#{runTreeLabel}' handler|]) >>= \case
      Left err -> return (Failure err, emptyExtraTimingInfo)
      Right x -> pure x
startTree node@(RunNodeDescribe {..}) ctx' = do
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    ((L.length . L.filter isFailure) <$> runNodesSequentially runNodeChildren ctx) >>= \case
      0 -> return (Success, emptyExtraTimingInfo)
      n -> return (Failure (ChildrenFailed Nothing n), emptyExtraTimingInfo)
startTree node@(RunNodeParallel {..}) ctx' = do
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    ((L.length . L.filter isFailure) <$> runNodesConcurrently runNodeCommon runNodeChildren ctx) >>= \case
      0 -> return (Success, emptyExtraTimingInfo)
      n -> return (Failure (ChildrenFailed Nothing n), emptyExtraTimingInfo)
startTree node@(RunNodeIt {..}) ctx' = do
  let RunNodeCommonWithStatus {..} = runNodeCommon
  let ctx = modifyBaseContext ctx' $ baseContextFromCommon runNodeCommon
  runInAsync node ctx $ do
    let label = runTreeLabel
    (results, _, _) <- timed runNodeCommon runTreeRecordTime (getBaseContext ctx) label $
      runExampleM runNodeCommon label runNodeExample ctx runTreeLogs Nothing
    return (results, emptyExtraTimingInfo)

-- * Util

runInAsync :: (HasBaseContext context, MonadIO m) => RunNode context -> context -> IO (Result, ExtraTimingInfo) -> m (Async Result)
runInAsync node ctx action = do
  let RunNodeCommonWithStatus {..} = runNodeCommon node
  let bc@(BaseContext {..}) = getBaseContext ctx
  let timerFn = if runTreeRecordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack runTreeLabel) else id
  let asyncName = T.pack [i|node #{runTreeId}, #{runTreeLabel}|]
  startTime <- liftIO getCurrentTime
  mvar <- liftIO newEmptyMVar
  myAsync <- liftIO $ managedAsyncWithUnmask baseContextRunId asyncName $ \unmask -> do
    flip withException (recordExceptionInStatus runTreeStatus) $ unmask $ do
      readMVar mvar
      (result, extraTimingInfo) <- timerFn action
      endTime <- liftIO getCurrentTime
      liftIO $ atomically $ writeTVar runTreeStatus $ Done startTime (setupFinishTime extraTimingInfo) (teardownStartTime extraTimingInfo) endTime result
      liftIO $ emitEvent baseContextOptions endTime runTreeId runTreeLabel (EventDone result)

      whenFailure result $ \reason -> do
        -- Make sure the folder exists, if configured
        whenJust baseContextPath $ createDirectoryIfMissing True

        -- Create error symlink when configured to
        case node of
          RunNodeDescribe {} -> return () -- These are just noisy so don't create them
          RunNodeParallel {} -> return () -- These are just noisy so don't create them
          _ -> do
            whenJust baseContextErrorSymlinksDir $ \errorsDir ->
              whenJust baseContextPath $ \dir -> do
#ifndef mingw32_HOST_OS
                whenJust baseContextRunRoot $ \runRoot -> do
#else
                whenJust baseContextRunRoot $ \_runRoot -> do
#endif
                  let symlinkBaseName = case runTreeLoc of
                        Nothing -> takeFileName dir
                        Just loc -> [i|#{srcLocFile loc}_line#{srcLocStartLine loc}_#{takeFileName dir}|]
                  let symlinkPath = errorsDir </> (nodeToFolderName symlinkBaseName 9999999 runTreeId)

                  -- Delete the symlink if it's already present. This can happen when re-running
                  -- a previously failed test
                  exists <- doesPathExist symlinkPath
                  when exists $ removePathForcibly symlinkPath

#ifndef mingw32_HOST_OS
                  -- Get a relative path from the error dir to the results dir. System.FilePath doesn't want to
                  -- introduce ".." components, so we have to do it ourselves
                  let errorDirDepth = L.length $ splitPath $ makeRelative runRoot errorsDir

                  -- Don't do createDirectoryLink on Windows, as creating symlinks is generally not allowed for users.
                  -- See https://security.stackexchange.com/questions/10194/why-do-you-have-to-be-an-admin-to-create-a-symlink-in-windows
                  -- TODO: could we detect if this permission is available?
                  let relativePath = joinPath (L.replicate errorDirDepth "..") </> (makeRelative runRoot dir)
                  liftIO $ createDirectoryLink relativePath symlinkPath
#endif

        -- Write failure info
        whenJust baseContextPath $ \dir -> do
          withFile (dir </> "failure.txt") AppendMode $ \h -> do
            -- Use the PrintFormatter to format failure.txt nicely
            let pf = defaultPrintFormatter {
                  printFormatterUseColor = False
                  , printFormatterLogLevel = Just LevelDebug
                  , printFormatterIncludeCallStacks = True
                  }
            flip runReaderT (pf, 0, h) $ do
              printFailureReason reason
              whenJust (failureCallStack reason) $ \cs -> do
                p "\n"
                printCallStack cs
              p "\n"
              printLogs runTreeLogs

      return result
  liftIO $ atomically $ writeTVar runTreeStatus $ Running startTime Nothing Nothing myAsync
  liftIO $ emitEvent baseContextOptions startTime runTreeId runTreeLabel EventStarted
  liftIO $ putMVar mvar ()
  return myAsync  -- TODO: fix race condition with writing to runTreeStatus (here and above)

-- | Run a list of children sequentially, cancelling everything on async exception
runNodesSequentially :: HasBaseContext context => [RunNode context] -> context -> IO [Result]
runNodesSequentially children ctx =
  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $
    forM (L.filter (shouldRunChild ctx) children) $ \child ->
      startTree child ctx >>= wait

-- | Run a list of children concurrently, cancelling everything on async exception
runNodesConcurrently :: forall context. HasBaseContext context => RunNodeCommon -> [RunNode context] -> context -> IO [Result]
runNodesConcurrently (RunNodeCommonWithStatus {runTreeLabel, runTreeId}) children ctx =
  flip withException (\(e :: SomeAsyncException) -> cancelAllChildrenWith children e) $
    mapM wait =<< sequence [startTree child (modifyTimingProfile ix ctx)
                           | (child, ix) <- L.zip runnableChildren [0..]]
  where
    runnableChildren = L.filter (shouldRunChild ctx) children

    leftPadWithZeros :: Int -> String
    leftPadWithZeros num = L.replicate (L.length (show (L.length runnableChildren)) - L.length (show num)) '0' <> show num

    modifyTimingProfile :: Int -> context -> context
    modifyTimingProfile n = flip modifyBaseContext (modifyTimingProfile' n)

    modifyTimingProfile' :: Int -> BaseContext -> BaseContext
    modifyTimingProfile' n bc@(BaseContext {..}) = bc {
      baseContextTestTimerProfile = baseContextTestTimerProfile <> [i|-#{runTreeLabel}-#{runTreeId}-#{leftPadWithZeros n}|]
      }

markAllChildrenWithResult :: (MonadIO m, HasBaseContext context') => [RunNode context] -> context' -> Result -> m ()
markAllChildrenWithResult children baseContext status = do
  now <- liftIO getCurrentTime
  forM_ (L.filter (shouldRunChild' baseContext) $ concatMap getCommons children) $ \child ->
    liftIO $ atomically $ modifyTVar' (runTreeStatus child) $ \case
      Running {..} -> Done now statusSetupFinishTime statusTeardownStartTime now status
      done@(Done {}) -> done { statusResult = status }
      _ -> Done now Nothing Nothing now status

cancelAllChildrenWith :: [RunNode context] -> SomeAsyncException -> IO ()
cancelAllChildrenWith children e = do
  forM_ children $ \node ->
    readTVarIO (runTreeStatus $ runNodeCommon node) >>= \case
      Running {..} -> cancelWith statusAsync e
      NotStarted -> do
        now <- getCurrentTime
        let reason = GotAsyncException Nothing Nothing (SomeAsyncExceptionWithEq e)
        atomically $ writeTVar (runTreeStatus $ runNodeCommon node) (Done now Nothing Nothing now (Failure reason))
      _ -> return ()

shouldRunChild :: (HasBaseContext ctx) => ctx -> RunNodeWithStatus context s l t -> Bool
shouldRunChild ctx node = shouldRunChild' ctx (runNodeCommon node)

shouldRunChild' :: (HasBaseContext ctx) => ctx -> RunNodeCommonWithStatus s l t -> Bool
shouldRunChild' ctx common = case baseContextOnlyRunIds $ getBaseContext ctx of
  Nothing -> True
  Just ids -> (runTreeId common) `S.member` ids

-- * Running examples

runExampleM :: (
  HasBaseContext r
  ) => RunNodeCommonWithStatus (TVar Status) l t -> String -> ExampleM r () -> r -> TVar (Seq LogEntry) -> Maybe String -> IO Result
runExampleM rnc label ex ctx logs exceptionMessage = runExampleM' rnc label ex ctx logs exceptionMessage >>= \case
  Left err -> return $ Failure err
  Right () -> return Success

runExampleM' :: (
  HasBaseContext r
  ) => RunNodeCommonWithStatus (TVar Status) l t -> String -> ExampleM r a -> r -> TVar (Seq LogEntry) -> Maybe String -> IO (Either FailureReason a)
runExampleM' rnc label ex ctx logs exceptionMessage = do
  maybeTestDirectory <- getTestDirectory ctx
  let options = baseContextOptions $ getBaseContext ctx

  -- We want our handleAny call to be *inside* the withLogFn call, because
  -- withFile will catch IOException and fill in its own information, making the
  -- resulting error confusing
  withLogFn maybeTestDirectory options $ \logFn ->
    handleAny (wrapInFailureReasonIfNecessary exceptionMessage) $
    withMaybeWarnOnLongExecution (getBaseContext ctx) rnc label $
    withMaybeCancelOnLongExecution (getBaseContext ctx) rnc label $
    (Right <$> (runLoggingT (runReaderT (unExampleT ex) ctx) logFn))

  where
    withLogFn :: Maybe FilePath -> Options -> (LogFn -> IO a) -> IO a
    withLogFn Nothing (Options {..}) action = action (withLateLogCheck optionsLateLogFile $ withBroadcast optionsLogBroadcast $ logToMemory optionsSavedLogLevel logs)
    withLogFn (Just logPath) (Options {..}) action = withFile (logPath </> "test_logs.txt") AppendMode $ \h -> do
      hSetBuffering h LineBuffering
      action (withLateLogCheck optionsLateLogFile $ withBroadcast optionsLogBroadcast $ logToMemoryAndFile optionsMemoryLogLevel optionsSavedLogLevel optionsLogFormatter logs h)

    withBroadcast :: Maybe (TChan (Int, String, LogEntry)) -> LogFn -> LogFn
    withBroadcast Nothing logFn = logFn
    withBroadcast (Just chan) logFn = \loc logSrc logLevel logStr -> do
      logFn loc logSrc logLevel logStr
      ts <- getCurrentTime
      atomically $ writeTChan chan (runTreeId rnc, runTreeLabel rnc, LogEntry ts loc logSrc logLevel (fromLogStr logStr))

    withLateLogCheck :: Maybe Handle -> LogFn -> LogFn
    withLateLogCheck Nothing logFn = logFn
    withLateLogCheck (Just lateH) logFn = \loc logSrc logLevel logStr -> do
      logFn loc logSrc logLevel logStr
      status <- readTVarIO (runTreeStatus rnc)
      case status of
        Done {} -> do
          ts <- getCurrentTime
          let bs = fromLogStr logStr
          let line = [i|#{show ts} [LATE] [#{runTreeId rnc}] #{runTreeLabel rnc}: #{BS8.unpack bs}\n|]
          BS8.hPutStr lateH (BS8.pack line)
          hFlush lateH
        _ -> return ()

    getTestDirectory :: (HasBaseContext a) => a -> IO (Maybe FilePath)
    getTestDirectory (getBaseContext -> (BaseContext {..})) = case baseContextPath of
      Nothing -> return Nothing
      Just dir -> do
        createDirectoryIfMissing True dir
        return $ Just dir

    wrapInFailureReasonIfNecessary :: Maybe String -> SomeException -> IO (Either FailureReason a)
    wrapInFailureReasonIfNecessary msg e = return $ Left $ case fromException e of
      Just (x :: FailureReason) -> x
      _ -> case fromException e of
        Just (SomeExceptionWithCallStack e' cs) -> GotException (Just cs) msg (SomeExceptionWithEq (SomeException e'))
        _ -> GotException Nothing msg (SomeExceptionWithEq e)

addSetupFinishTimeToStatus :: (MonadIO m) => TVar Status -> UTCTime -> m ()
addSetupFinishTimeToStatus statusVar setupFinishTime = atomically $ modifyTVar' statusVar $ \case
  status@(Running {}) -> status { statusSetupFinishTime = Just setupFinishTime }
  status@(Done {}) -> status { statusSetupFinishTime = Just setupFinishTime }
  x -> x

addTeardownStartTimeToStatus :: (MonadIO m) => TVar Status -> UTCTime -> m ()
addTeardownStartTimeToStatus statusVar t = atomically $ modifyTVar' statusVar $ \case
  status@(Running {}) -> status { statusTeardownStartTime = Just t }
  status@(Done {}) -> status { statusTeardownStartTime = Just t }
  x -> x

emitEvent :: Options -> UTCTime -> Int -> String -> NodeEventType -> IO ()
emitEvent (Options {optionsEventBroadcast = Just chan}) ts nid label evtType =
  atomically $ writeTChan chan (NodeEvent ts nid label evtType)
emitEvent _ _ _ _ _ = return ()

recordExceptionInStatus :: (MonadIO m) => TVar Status -> SomeException -> m ()
recordExceptionInStatus status e = do
  endTime <- liftIO getCurrentTime
  let ret = case fromException e of
        Just (e' :: SomeAsyncException) -> Failure (GotAsyncException Nothing Nothing (SomeAsyncExceptionWithEq e'))
        _ -> case fromException e of
          Just (e' :: FailureReason) -> Failure e'
          _ -> Failure (GotException Nothing Nothing (SomeExceptionWithEq e))
  liftIO $ atomically $ modifyTVar' status $ \case
    Running {..} -> Done statusStartTime statusSetupFinishTime statusTeardownStartTime endTime ret
    _ -> Done endTime Nothing Nothing endTime ret

timed :: forall m a s l t. (MonadUnliftIO m) => RunNodeCommonWithStatus s l t -> Bool -> BaseContext -> String -> m a -> m (a, UTCTime, UTCTime)
timed _rnc recordTime bc@(BaseContext {..}) label action = do
  let timerFn = if recordTime then timeAction' (getTestTimer bc) baseContextTestTimerProfile (T.pack label) else id

  startTime <- liftIO getCurrentTime
  ret <- timerFn action
  endTime <- liftIO getCurrentTime
  pure (ret, startTime, endTime)

withMaybeWarnOnLongExecution :: (MonadUnliftIO m) => BaseContext -> RunNodeCommonWithStatus s l t -> String -> m a -> m a
withMaybeWarnOnLongExecution (BaseContext {..}) rnc@(RunNodeCommonWithStatus {..}) label inner = case optionsWarnOnLongExecutionMs baseContextOptions of
  Nothing -> inner
  Just maxTimeMs -> do
    managedWithAsync_ baseContextRunId (T.pack [i|warn-long:#{runTreeId}:#{label}|]) (waiter maxTimeMs) inner

  where
    waiter maxTimeMs = do
      -- Use unsafeUnmask to make threadDelay interruptible even in masked cleanup contexts
      -- (like the bracket call of an Introduce node)
      liftIO $ unsafeUnmask $ threadDelay (maxTimeMs * 1000)

      micros <- liftIO getCurrentMicrosInteger

      whenJust baseContextRunRoot $ \runRoot -> do
        let dir = runRoot </> "long-execution-warnings"
        createDirectoryIfMissing True dir
        writeInfoToFile rnc maxTimeMs (dir </> (nodeToFolderName [i|#{micros}-#{runTreeId}-warn-#{label}|] 1 0) <.> "log") label

withMaybeCancelOnLongExecution :: (MonadUnliftIO m) => BaseContext -> RunNodeCommonWithStatus s l t -> String -> m a -> m a
withMaybeCancelOnLongExecution (BaseContext {..}) rnc@(RunNodeCommonWithStatus {..}) label inner = case optionsCancelOnLongExecutionMs baseContextOptions of
  Nothing -> inner
  Just maxTimeMs ->
    race (waiter maxTimeMs) inner >>= \case
      Left _ -> do
        micros <- liftIO getCurrentMicrosInteger
        whenJust baseContextRunRoot $ \runRoot -> do
          let dir = runRoot </> "long-execution-warnings"
          createDirectoryIfMissing True dir
          writeInfoToFile rnc maxTimeMs (dir </> (nodeToFolderName [i|#{micros}-#{runTreeId}-cancel-#{label}|] 1 0) <.> "log") label
        throwIO $ Reason (Just callStack) $ [i|Timed out long-running node|]
      Right x -> return x
  where
    waiter maxTimeMs =
      -- Use unsafeUnmask to make threadDelay interruptible even in masked cleanup contexts
      liftIO $ unsafeUnmask $ threadDelay (maxTimeMs * 1000)

writeInfoToFile :: MonadUnliftIO m => RunNodeCommonWithStatus s l t -> Int -> FilePath -> String -> m ()
writeInfoToFile (RunNodeCommonWithStatus {..}) ms fileName label =
  bracket (liftIO $ openFile fileName WriteMode) (liftIO . hClose) $ \h -> do
    liftIO $ T.hPutStrLn h [__i|Node was still running after #{ms}ms: #{label}

                                Ancestors: #{runTreeAncestors}

                                Folder: #{runTreeFolder}
                                |]
    whenJust runTreeLoc $ \(SrcLoc {..}) ->
      liftIO $ T.hPutStrLn h [__i|\nLoc: #{srcLocFile}:#{srcLocStartLine}:#{srcLocStartCol} in #{srcLocPackage}:#{srcLocModule}
                                 |]

getCurrentMicrosInteger :: IO Integer
getCurrentMicrosInteger = do
  t <- getPOSIXTime
  return $ round (t * 1000000)
