{-# LANGUAGE RecordWildCards #-}
module System.Console.AsciiProgress
    ( ProgressBar(..)
    , Options(..)
    , Stats(..)
    , isComplete
    , newProgressBar
    , complete
    , tick
    , tickN
    , tickNI
    , getProgressStrIO
    , getProgressStats
    , getProgressStr
    -- Re-exports:
    , Default(..)
    , module System.Console.Regions
    )
  where

import           Control.Applicative                   ((<$>))
import           Control.Concurrent                    (modifyMVar_, readChan,
                                                        readMVar, writeChan)
import           Control.Concurrent.Async              (Async (..), async, poll,
                                                        wait)
import           Data.Default                          (Default (..))
import           Data.Maybe                            (fromMaybe, isJust)
import           System.Console.AsciiProgress.Internal
import           System.Console.Regions

data ProgressBar = ProgressBar { ProgressBar -> ProgressBarInfo
pgInfo   :: ProgressBarInfo
                               , ProgressBar -> Async ()
pgFuture :: Async ()
                               , ProgressBar -> ConsoleRegion
pgRegion :: ConsoleRegion
                               }

-- |
-- Creates a new progress bar with the given @Options@. Multiple progress bars
-- may be created. This package depends on `concurrent-output`, so it's --
-- necessary that progress-bar usage is wrapped with a call to
-- 'displayConsoleRegions'.
--
-- > import           Control.Concurrent           (threadDelay)
-- > import           Control.Monad                (unless)
-- > import           System.Console.AsciiProgress
-- >
-- > main :: IO ()
-- > main = displayConsoleRegions $ do
-- >    pg <- newProgressBar def { pgWidth = 100
-- >                             , pgOnCompletion = Just "Done :percent after :elapsed seconds"
-- >                             }
-- >    loop pg
-- >  where
-- >    loop pg = do
-- >        b <- isComplete pg
-- >        unless b $ do
-- >            threadDelay $ 200 * 1000
-- >            tick pg
-- >            loop pg
newProgressBar :: Options -> IO ProgressBar
newProgressBar :: Options -> IO ProgressBar
newProgressBar Options
opts = do
    ConsoleRegion
region <- RegionLayout -> IO ConsoleRegion
forall (m :: * -> *).
LiftRegion m =>
RegionLayout -> m ConsoleRegion
openConsoleRegion RegionLayout
Linear
    ProgressBarInfo
info <- Options -> IO ProgressBarInfo
newProgressBarInfo Options
opts

    -- Display initial progress-bar
    String
pgStr <- Options -> Options -> Stats -> String
pgGetProgressStr Options
opts Options
opts (Stats -> String) -> IO Stats -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProgressBarInfo -> IO Stats
getInfoStats ProgressBarInfo
info
    ConsoleRegion -> String -> IO ()
forall v (m :: * -> *).
(ToRegionContent v, LiftRegion m) =>
ConsoleRegion -> v -> m ()
setConsoleRegion ConsoleRegion
region String
pgStr

    Async ()
future <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ ProgressBarInfo -> ConsoleRegion -> IO ()
start ProgressBarInfo
info ConsoleRegion
region
    ProgressBar -> IO ProgressBar
forall (m :: * -> *) a. Monad m => a -> m a
return (ProgressBar -> IO ProgressBar) -> ProgressBar -> IO ProgressBar
forall a b. (a -> b) -> a -> b
$ ProgressBarInfo -> Async () -> ConsoleRegion -> ProgressBar
ProgressBar ProgressBarInfo
info Async ()
future ConsoleRegion
region
  where
    start :: ProgressBarInfo -> ConsoleRegion -> IO ()
start info :: ProgressBarInfo
info@ProgressBarInfo{Chan Integer
MVar Integer
MVar UTCTime
Options
pgFirstTick :: ProgressBarInfo -> MVar UTCTime
pgCompleted :: ProgressBarInfo -> MVar Integer
pgChannel :: ProgressBarInfo -> Chan Integer
pgOptions :: ProgressBarInfo -> Options
pgFirstTick :: MVar UTCTime
pgCompleted :: MVar Integer
pgChannel :: Chan Integer
pgOptions :: Options
..} ConsoleRegion
region = do
       Integer
c <- MVar Integer -> IO Integer
forall a. MVar a -> IO a
readMVar MVar Integer
pgCompleted
       Integer -> IO () -> IO ()
unlessDone Integer
c (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
           Integer
n <- Chan Integer -> IO Integer
forall a. Chan a -> IO a
readChan Chan Integer
pgChannel
           ()
_ <- ProgressBarInfo -> ConsoleRegion -> Integer -> IO ()
handleMessage ProgressBarInfo
info ConsoleRegion
region Integer
n
           Integer -> IO () -> IO ()
unlessDone (Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ ProgressBarInfo -> ConsoleRegion -> IO ()
start ProgressBarInfo
info ConsoleRegion
region
      where
        unlessDone :: Integer -> IO () -> IO ()
unlessDone Integer
c IO ()
action | Integer
c Integer -> Integer -> Bool
forall a. Ord a => a -> a -> Bool
< Options -> Integer
pgTotal Options
opts = IO ()
action
        unlessDone Integer
_ IO ()
_ = do
            let fmt :: String
fmt = String -> Maybe String -> String
forall a. a -> Maybe a -> a
fromMaybe (Options -> String
pgFormat Options
opts) (Options -> Maybe String
pgOnCompletion Options
opts)
            String
onCompletion <- Options -> Options -> Stats -> String
pgGetProgressStr Options
opts Options
opts { pgFormat :: String
pgFormat = String
fmt } (Stats -> String) -> IO Stats -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProgressBarInfo -> IO Stats
getInfoStats ProgressBarInfo
info
            ConsoleRegion -> String -> IO ()
forall v (m :: * -> *).
(ToRegionContent v, LiftRegion m) =>
ConsoleRegion -> v -> m ()
setConsoleRegion ConsoleRegion
region String
onCompletion

    handleMessage :: ProgressBarInfo -> ConsoleRegion -> Integer -> IO ()
handleMessage ProgressBarInfo
info ConsoleRegion
region Integer
n = do
        -- Update the completed tick count
        MVar Integer -> (Integer -> IO Integer) -> IO ()
forall a. MVar a -> (a -> IO a) -> IO ()
modifyMVar_ (ProgressBarInfo -> MVar Integer
pgCompleted ProgressBarInfo
info) (\Integer
c -> Integer -> IO Integer
forall (m :: * -> *) a. Monad m => a -> m a
return (Integer
c Integer -> Integer -> Integer
forall a. Num a => a -> a -> a
+ Integer
n))
        -- Find and update the current and first tick times:
        Stats
stats <- ProgressBarInfo -> IO Stats
getInfoStats ProgressBarInfo
info
        let progressStr :: String
progressStr = Options -> Options -> Stats -> String
pgGetProgressStr Options
opts Options
opts Stats
stats
        ConsoleRegion -> String -> IO ()
forall v (m :: * -> *).
(ToRegionContent v, LiftRegion m) =>
ConsoleRegion -> v -> m ()
setConsoleRegion ConsoleRegion
region String
progressStr

-- |
-- Tick the progress bar
tick :: ProgressBar -> IO ()
tick :: ProgressBar -> IO ()
tick ProgressBar
pg = ProgressBar -> Int -> IO ()
tickN ProgressBar
pg Int
1

-- |
-- Tick the progress bar N times
tickN :: ProgressBar -> Int -> IO ()
tickN :: ProgressBar -> Int -> IO ()
tickN (ProgressBar ProgressBarInfo
info Async ()
_ ConsoleRegion
_) = Chan Integer -> Integer -> IO ()
forall a. Chan a -> a -> IO ()
writeChan (ProgressBarInfo -> Chan Integer
pgChannel ProgressBarInfo
info) (Integer -> IO ()) -> (Int -> Integer) -> Int -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral

-- |
-- Tick the progress bar N times
tickNI :: ProgressBar -> Integer -> IO ()
tickNI :: ProgressBar -> Integer -> IO ()
tickNI (ProgressBar ProgressBarInfo
info Async ()
_ ConsoleRegion
_) = Chan Integer -> Integer -> IO ()
forall a. Chan a -> a -> IO ()
writeChan (ProgressBarInfo -> Chan Integer
pgChannel ProgressBarInfo
info)

-- |
-- Returns if the progress bar rendering thread has exited (it has done enough
-- ticks)
isComplete :: ProgressBar -> IO Bool
isComplete :: ProgressBar -> IO Bool
isComplete (ProgressBar ProgressBarInfo
_ Async ()
future ConsoleRegion
_) = Maybe (Either SomeException ()) -> Bool
forall a. Maybe a -> Bool
isJust (Maybe (Either SomeException ()) -> Bool)
-> IO (Maybe (Either SomeException ())) -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Async () -> IO (Maybe (Either SomeException ()))
forall a. Async a -> IO (Maybe (Either SomeException a))
poll Async ()
future

-- |
-- Forces a 'ProgressBar' to finish
complete :: ProgressBar -> IO ()
complete :: ProgressBar -> IO ()
complete pg :: ProgressBar
pg@(ProgressBar ProgressBarInfo
info Async ()
future ConsoleRegion
_) = do
    let total :: Integer
total = Options -> Integer
pgTotal (ProgressBarInfo -> Options
pgOptions ProgressBarInfo
info)
    ProgressBar -> Integer -> IO ()
tickNI ProgressBar
pg Integer
total
    Async () -> IO ()
forall a. Async a -> IO a
wait Async ()
future

-- |
-- Gets the progress bar current @Stats @object
getProgressStats :: ProgressBar -> IO Stats
getProgressStats :: ProgressBar -> IO Stats
getProgressStats (ProgressBar ProgressBarInfo
info Async ()
_ ConsoleRegion
_) = ProgressBarInfo -> IO Stats
getInfoStats ProgressBarInfo
info

-- |
-- Like @getProgressStr@ but works on the @ProgressBar@ object and uses the IO
-- monad.
getProgressStrIO :: ProgressBar -> IO String
getProgressStrIO :: ProgressBar -> IO String
getProgressStrIO (ProgressBar ProgressBarInfo
info Async ()
_ ConsoleRegion
_) =
    Options -> Stats -> String
getProgressStr (ProgressBarInfo -> Options
pgOptions ProgressBarInfo
info) (Stats -> String) -> IO Stats -> IO String
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ProgressBarInfo -> IO Stats
getInfoStats ProgressBarInfo
info