{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE NoFieldSelectors #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

module Skeletest.Internal.Utils.Term (
  init,
  setANSISupport,

  -- * Global attributes
  Handle,
  width,
  stdout,
  stderr,
  supportsANSI,

  -- * Output helpers
  flush,
  output,
  outputN,
  outputErr,
  outputErrN,
  clearInPlace,
  outputInPlace,
) where

import Control.Exception (evaluate)
import Control.Monad (forM_, replicateM_)
import Data.Foldable (traverse_)
import Data.String.AnsiEscapeCodes.Strip.Text (stripAnsiEscapeCodes)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.IO qualified as Text
import GHC.IO.Handle qualified as IO
import System.Console.ANSI qualified as ANSI
import System.Console.Terminal.Size qualified as TermSize
import System.IO qualified as IO
import System.IO.Unsafe (unsafePerformIO)
import UnliftIO.MVar (MVar, modifyMVar_, newMVar, readMVar, withMVar)
import Prelude hiding (init)

data GlobalTermData = GlobalTermData
  { GlobalTermData -> Int
width :: Int
  , GlobalTermData -> Handle
stdout :: Handle
  , GlobalTermData -> Handle
stderr :: Handle
  }

newtype Handle = Handle {Handle -> MVar Handle'
var :: MVar Handle'}

data Handle' = Handle'
  { Handle' -> Handle
handle :: IO.Handle
  , Handle' -> Bool
supportsANSI :: Bool
  , Handle' -> Maybe Text
lastInPlaceOutput :: Maybe Text
  -- ^ See 'outputInPlace'
  }

globalTermData :: GlobalTermData
globalTermData :: GlobalTermData
globalTermData = IO GlobalTermData -> GlobalTermData
forall a. IO a -> a
unsafePerformIO (IO GlobalTermData -> GlobalTermData)
-> IO GlobalTermData -> GlobalTermData
forall a b. (a -> b) -> a -> b
$ do
  Int
width <- Int -> (Window Int -> Int) -> Maybe (Window Int) -> Int
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Int
80 Window Int -> Int
forall a. Window a -> a
TermSize.width (Maybe (Window Int) -> Int) -> IO (Maybe (Window Int)) -> IO Int
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO (Maybe (Window Int))
forall n. Integral n => IO (Maybe (Window n))
TermSize.size
  Handle
stdout <- Handle -> IO Handle
getHandle Handle
IO.stdout
  Handle
stderr <- Handle -> IO Handle
getHandle Handle
IO.stderr
  GlobalTermData -> IO GlobalTermData
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure GlobalTermData{Int
Handle
width :: Int
stdout :: Handle
stderr :: Handle
width :: Int
stdout :: Handle
stderr :: Handle
..}
 where
  getHandle :: Handle -> IO Handle
getHandle Handle
h = do
    Handle
handle <- Handle -> IO Handle
IO.hDuplicate Handle
h
    Bool
supportsANSI <- Handle -> IO Bool
ANSI.hSupportsANSI Handle
handle
    let lastInPlaceOutput :: Maybe a
lastInPlaceOutput = Maybe a
forall a. Maybe a
Nothing
    MVar Handle' -> Handle
Handle (MVar Handle' -> Handle) -> IO (MVar Handle') -> IO Handle
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Handle' -> IO (MVar Handle')
forall (m :: * -> *) a. MonadIO m => a -> m (MVar a)
newMVar Handle'{Bool
Maybe Text
Handle
forall a. Maybe a
handle :: Handle
supportsANSI :: Bool
lastInPlaceOutput :: Maybe Text
handle :: Handle
supportsANSI :: Bool
lastInPlaceOutput :: forall a. Maybe a
..}
{-# NOINLINE globalTermData #-}

init :: IO ()
init :: IO ()
init = do
  -- Configure stdout/stderr globally, both for Term.output and for anything
  -- writing directly to stdout/stderr (e.g. user tests with --capture-output=off)
  [Handle] -> (Handle -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Handle
IO.stdout, Handle
IO.stderr] ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle
h -> do
    Handle -> TextEncoding -> IO ()
IO.hSetEncoding Handle
h TextEncoding
IO.utf8
    Handle -> BufferMode -> IO ()
IO.hSetBuffering Handle
h BufferMode
IO.LineBuffering

  -- Make sure globalTermData is initialized
  GlobalTermData
_ <- GlobalTermData -> IO GlobalTermData
forall a. a -> IO a
evaluate GlobalTermData
globalTermData
  () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

-- Use the terminal width at the beginning of the test suite; don't
-- handle users changing terminal width in the middle right now
width :: Int
width :: Int
width = GlobalTermData
globalTermData.width

stdout :: Handle
stdout :: Handle
stdout = GlobalTermData
globalTermData.stdout

stderr :: Handle
stderr :: Handle
stderr = GlobalTermData
globalTermData.stderr

supportsANSI :: Handle -> IO Bool
supportsANSI :: Handle -> IO Bool
supportsANSI (Handle MVar Handle'
hvar) = (.supportsANSI) (Handle' -> Bool) -> IO Handle' -> IO Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MVar Handle' -> IO Handle'
forall (m :: * -> *) a. MonadIO m => MVar a -> m a
readMVar MVar Handle'
hvar

setANSISupport :: Bool -> IO ()
setANSISupport :: Bool -> IO ()
setANSISupport Bool
x =
  [Handle] -> (Handle -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [GlobalTermData
globalTermData.stdout, GlobalTermData
globalTermData.stderr] ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Handle MVar Handle'
hvar) -> do
    MVar Handle' -> (Handle' -> IO Handle') -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
MVar a -> (a -> m a) -> m ()
modifyMVar_ MVar Handle'
hvar ((Handle' -> IO Handle') -> IO ())
-> (Handle' -> IO Handle') -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h -> Handle' -> IO Handle'
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle'
h{supportsANSI = x}

flush :: IO ()
flush :: IO ()
flush = do
  [Handle] -> (Handle -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [GlobalTermData
globalTermData.stdout, GlobalTermData
globalTermData.stderr] ((Handle -> IO ()) -> IO ()) -> (Handle -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \(Handle MVar Handle'
hvar) -> do
    MVar Handle' -> (Handle' -> IO ()) -> IO ()
forall (m :: * -> *) a b.
MonadUnliftIO m =>
MVar a -> (a -> m b) -> m b
withMVar MVar Handle'
hvar ((Handle' -> IO ()) -> IO ()) -> (Handle' -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h -> Handle -> IO ()
IO.hFlush Handle'
h.handle

output :: Text -> IO ()
output :: Text -> IO ()
output = Handle -> Text -> IO ()
outputWith GlobalTermData
globalTermData.stdout

outputN :: Text -> IO ()
outputN :: Text -> IO ()
outputN = Handle -> Text -> IO ()
outputNWith GlobalTermData
globalTermData.stdout

outputErr :: Text -> IO ()
outputErr :: Text -> IO ()
outputErr = Handle -> Text -> IO ()
outputWith GlobalTermData
globalTermData.stderr

outputErrN :: Text -> IO ()
outputErrN :: Text -> IO ()
outputErrN = Handle -> Text -> IO ()
outputNWith GlobalTermData
globalTermData.stderr

outputWith :: Handle -> Text -> IO ()
outputWith :: Handle -> Text -> IO ()
outputWith (Handle MVar Handle'
hvar) Text
s = MVar Handle' -> (Handle' -> IO Handle') -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
MVar a -> (a -> m a) -> m ()
modifyMVar_ MVar Handle'
hvar ((Handle' -> IO Handle') -> IO ())
-> (Handle' -> IO Handle') -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h -> Handle' -> Text -> IO Handle'
outputWith' Handle'
h Text
s

outputNWith :: Handle -> Text -> IO ()
outputNWith :: Handle -> Text -> IO ()
outputNWith (Handle MVar Handle'
hvar) Text
s = MVar Handle' -> (Handle' -> IO Handle') -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
MVar a -> (a -> m a) -> m ()
modifyMVar_ MVar Handle'
hvar ((Handle' -> IO Handle') -> IO ())
-> (Handle' -> IO Handle') -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h -> Handle' -> Text -> IO Handle'
outputNWith' Handle'
h Text
s

outputWith' :: Handle' -> Text -> IO Handle'
outputWith' :: Handle' -> Text -> IO Handle'
outputWith' Handle'
h Text
s = do
  Handle -> Text -> IO ()
Text.hPutStrLn Handle'
h.handle Text
s
  Handle' -> IO Handle'
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle'
h{lastInPlaceOutput = Nothing}

outputNWith' :: Handle' -> Text -> IO Handle'
outputNWith' :: Handle' -> Text -> IO Handle'
outputNWith' Handle'
h Text
s = do
  Handle -> Text -> IO ()
Text.hPutStr Handle'
h.handle Text
s
  Handle -> IO ()
IO.hFlush Handle'
h.handle
  Handle' -> IO Handle'
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle'
h{lastInPlaceOutput = Nothing}

-- | Same as 'outputN', except if called multiple times in a row, will update
-- the message in-place.
--
-- Should only be used when stdout supports ANSI, but this isn't checked.
--
-- Long-term, should probably be replaced with concurrent-output, but it's
-- currently hardcoded to IO.stdout.
outputInPlace :: Text -> IO ()
outputInPlace :: Text -> IO ()
outputInPlace Text
s = MVar Handle' -> (Handle' -> IO Handle') -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
MVar a -> (a -> m a) -> m ()
modifyMVar_ GlobalTermData
globalTermData.stdout.var ((Handle' -> IO Handle') -> IO ())
-> (Handle' -> IO Handle') -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h0 -> do
  Handle'
h1 <- Handle' -> IO Handle'
clearInPlaceWith' Handle'
h0
  Handle'
h2 <- Handle' -> Text -> IO Handle'
outputNWith' Handle'
h1 Text
s
  Handle' -> IO Handle'
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle'
h2{lastInPlaceOutput = Just s}

-- | Clear last in-place output. See 'outputInPlace'.
clearInPlace :: IO ()
clearInPlace :: IO ()
clearInPlace = MVar Handle' -> (Handle' -> IO Handle') -> IO ()
forall (m :: * -> *) a.
MonadUnliftIO m =>
MVar a -> (a -> m a) -> m ()
modifyMVar_ GlobalTermData
globalTermData.stdout.var ((Handle' -> IO Handle') -> IO ())
-> (Handle' -> IO Handle') -> IO ()
forall a b. (a -> b) -> a -> b
$ \Handle'
h -> Handle' -> IO Handle'
clearInPlaceWith' Handle'
h

clearInPlaceWith' :: Handle' -> IO Handle'
clearInPlaceWith' :: Handle' -> IO Handle'
clearInPlaceWith' Handle'
h = do
  (Text -> IO ()) -> Maybe Text -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Text -> IO ()
clearLastInPlaceOutput Handle'
h.lastInPlaceOutput
  Handle' -> IO Handle'
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Handle'
h{lastInPlaceOutput = Nothing}
 where
  clearLastInPlaceOutput :: Text -> IO ()
clearLastInPlaceOutput Text
lastOutput = do
    let n :: Int
n = Text -> Int
calculateNumLines Text
lastOutput
    Handle -> Text -> IO ()
Text.hPutStr Handle'
h.handle Text
"\r\ESC[K" -- Clear the current line
    Int -> IO () -> IO ()
forall (m :: * -> *) a. Applicative m => Int -> m a -> m ()
replicateM_ (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
      -- Clear any lines above
      Handle -> Text -> IO ()
Text.hPutStr Handle'
h.handle Text
"\ESC[F\ESC[K"

  calculateNumLines :: Text -> Int
calculateNumLines =
    let calc :: Text -> Int
calc = [Text] -> Int
forall a. [a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
length ([Text] -> Int) -> (Text -> [Text]) -> Text -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Int -> Text -> [Text]
Text.chunksOf GlobalTermData
globalTermData.width (Text -> [Text]) -> (Text -> Text) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
stripAnsiEscapeCodes
     in [Int] -> Int
forall a. Num a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Num a) => t a -> a
sum ([Int] -> Int) -> (Text -> [Int]) -> Text -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Int) -> [Text] -> [Int]
forall a b. (a -> b) -> [a] -> [b]
map Text -> Int
calc ([Text] -> [Int]) -> (Text -> [Text]) -> Text -> [Int]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [Text]
Text.lines