{-# LANGUAGE RecordWildCards #-}

module General.Timing(Timing, withTiming, timed, timedOverwrite) where

import Data.List.Extra
import System.Time.Extra
import Data.IORef
import Control.Monad.Extra
import System.IO
import General.Util
import Control.Monad.IO.Class

-- | A mutable object to keep timing information
data Timing = Timing
    { -- | Get time since the initialization of this 'Timing'.
      timingOffset :: IO Seconds
      -- | Record timings for writing to a file
    , timingStore :: IORef [(String, Seconds)]
      -- | If you are below T you may overwrite N characters
      -- at the end of the current terminal output.
      -- Only used iff @timingTerminal == True@.
    , timingOverwrite :: IORef (Maybe (Seconds, Int))
      -- | whether is this a terminal
    , timingTerminal :: Bool
    }

-- | Time an action, printing timing information to the terminal
withTiming ::
  -- | A file to optionally write all timings to, after the action is finished
  Maybe FilePath ->
  -- | An action that can write timings into 'Timing'
  (Timing -> IO a) ->
  IO a
withTiming writeTimingsTo act = do
    timingOffset <- offsetTime
    timingStore <- newIORef []
    timingOverwrite <- newIORef Nothing
    timingTerminal <- hIsTerminalDevice stdout

    res <- act Timing{..}
    total <- timingOffset
    whenJust writeTimingsTo $ \file -> do
        xs <- readIORef timingStore
        -- Expecting unrecorded of ~2s
        -- Most of that comes from the pipeline - we get occasional 0.01 between items as one flushes
        -- Then at the end there is ~0.5 while the final item flushes
        xs <- pure $ sortOn (negate . snd) $ ("Unrecorded", total - sum (map snd xs)) : xs
        writeFile file $ unlines $ prettyTable 2 "Secs" xs
    putStrLn $ "Took " ++ showDuration total
    pure res


-- skip it if have written out in the last 1s and takes < 0.1

-- | Time & write the given message to stdout
timed :: MonadIO m => Timing -> String -> m a -> m a
timed = timedEx False

-- | Time & write the given message to stdout
-- overwriting a previous message if it was marked as overwritable
timedOverwrite :: MonadIO m => Timing -> String -> m a -> m a
timedOverwrite = timedEx True

timedEx :: MonadIO m => Bool -> Timing -> String -> m a -> m a
timedEx overwrite Timing{..} msg act = do
    start <- liftIO timingOffset
    liftIO $ whenJustM (readIORef timingOverwrite) $ \(t,n) ->
        if overwrite && start < t then
            putStr $ replicate n '\b' ++ replicate n ' ' ++ replicate n '\b'
        else
            putStrLn ""

    let out msg = liftIO $ putStr msg >> pure (length msg)
    undo1 <- out $ msg ++ "... "
    liftIO $ hFlush stdout

    res <- act
    end <- liftIO timingOffset
    let time = end - start
    liftIO $ modifyIORef timingStore ((msg,time):)

    s <- maybe "" (\x -> " (" ++ x ++ ")") <$> liftIO getStatsPeakAllocBytes
    undo2 <- out $ showDuration time ++ s

    old <- liftIO $ readIORef timingOverwrite
    let next = maybe (start + 1.0) fst old
    liftIO $ if timingTerminal && overwrite && end < next then
        writeIORef timingOverwrite $ Just (next, undo1 + undo2)
     else do
        writeIORef timingOverwrite Nothing
        putStrLn ""
    pure res