{-# LANGUAGE ViewPatterns #-}

-- | Provides utilities.
module Shrun.Utils
  ( -- * Text Utils
    breakStripPoint,
    truncateIfNeeded,
    stripControlAll,
    stripControlSmart,
    escapeDoubleQuotes,

    -- * MonadTime Utils
    diffTime,
    timeSpecToRelTime,
    foldMap1,

    -- * Misc Utils
    parseByteText,
    whileM_,
    whenLeft,
    untilJust,
    unsafeListToNESeq,
    (∸),
    readStripUnderscores,
  )
where

import Data.Bytes (Conversion (convert), SomeSize, parse)
import Data.Char (isControl, isLetter)
import Data.Either (either)
import Data.Sequence.NonEmpty qualified as NESeq
import Data.Text qualified as T
import Data.Text.Lazy qualified as TL
import Data.Text.Lazy.Builder (Builder)
import Data.Text.Lazy.Builder qualified as TLB
import Data.Time.Relative (RelativeTime, fromSeconds)
import Effects.Time (TimeSpec, diffTimeSpec)
import GHC.Exts (IsList (fromList))
import Shrun.Data.Text (UnlinedText)
import Shrun.Data.Text qualified as ShrunText
import Shrun.Prelude
import Text.Read (Read)
import Text.Read qualified as TR

-- $setup
-- >>> :set -XOverloadedLists
-- >>> :set -XTemplateHaskell
-- >>> import Data.List.NonEmpty (NonEmpty (..))
-- >>> import Data.Semigroup (Sum (..))
-- >>> import Data.Text qualified as T
-- >>> import Effects.Time (TimeSpec (..))
-- >>> import Shrun.Prelude

-- | For given \(x, y\), returns the absolute difference \(|x - y|\)
-- in seconds.
--
-- ==== __Examples__
-- >>> :{
--   let t1 = MkTimeSpec 5 0
--       -- 20 s + 1 billion ns = 21 s
--       t2 = MkTimeSpec 20 1_000_000_000
--   in diffTime t1 t2
-- :}
-- 16
diffTime :: TimeSpec -> TimeSpec -> Natural
diffTime :: TimeSpec -> TimeSpec -> Natural
diffTime TimeSpec
t1 TimeSpec
t2 = Optic' A_Lens NoIx TimeSpec Natural -> TimeSpec -> Natural
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx TimeSpec Natural
#sec (TimeSpec -> Natural) -> TimeSpec -> Natural
forall a b. (a -> b) -> a -> b
$ TimeSpec -> TimeSpec -> TimeSpec
diffTimeSpec TimeSpec
t1 TimeSpec
t2

-- | Transforms a 'TimeSpec' into a 'RelativeTime'.
timeSpecToRelTime :: TimeSpec -> RelativeTime
timeSpecToRelTime :: TimeSpec -> RelativeTime
timeSpecToRelTime = Natural -> RelativeTime
fromSeconds (Natural -> RelativeTime)
-> (TimeSpec -> Natural) -> TimeSpec -> RelativeTime
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Optic' A_Lens NoIx TimeSpec Natural -> TimeSpec -> Natural
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx TimeSpec Natural
#sec

-- | Relaxes 'foldMap'\'s 'Monoid' constraint to 'Semigroup'. Requires a
-- starting value. This will have to do until semigroupoids' Foldable1 is
-- in base.
--
-- ==== __Examples__
-- >>> foldMap1 @List Sum 0 [1..4]
-- Sum {getSum = 10}
--
-- >>> -- Silly, but demonstrates usage i.e. with non-monoid NonEmpty.
-- >>> foldMap1 @List (:| []) 1 [2,3,4]
-- 1 :| [2,3,4]
foldMap1 :: (Foldable f, Semigroup s) => (a -> s) -> a -> f a -> s
foldMap1 :: forall (f :: Type -> Type) s a.
(Foldable f, Semigroup s) =>
(a -> s) -> a -> f a -> s
foldMap1 a -> s
f a
x f a
xs = (a -> (a -> s) -> a -> s) -> (a -> s) -> f a -> a -> s
forall a b. (a -> b -> b) -> b -> f a -> b
forall (t :: Type -> Type) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\a
b a -> s
g a
y -> a -> s
f a
y s -> s -> s
forall a. Semigroup a => a -> a -> a
<> a -> s
g a
b) a -> s
f f a
xs a
x

-- | Wrapper for 'Text'\'s 'T.breakOn' that differs in that:
--
-- 1. If the @needle@ is found within the @haystack@, we do not include it
-- in the second part of the pair.
--
-- ==== __Examples__
-- >>> -- Data.Text
-- >>> T.breakOn "=" "HEY=LISTEN"
-- ("HEY","=LISTEN")
--
-- >>> -- Shrun.Utils.Text
-- >>> breakStripPoint "=" "HEY=LISTEN"
-- ("HEY","LISTEN")
--
-- Other examples:
--
-- >>> breakStripPoint "=" "HEYLISTEN"
-- ("HEYLISTEN","")
--
-- >>> breakStripPoint "=" "=HEYLISTEN"
-- ("","HEYLISTEN")
--
-- >>> breakStripPoint "=" "HEYLISTEN="
-- ("HEYLISTEN","")
--
-- >>> breakStripPoint "=" "HEY==LISTEN"
-- ("HEY","=LISTEN")
breakStripPoint :: Text -> Text -> Tuple2 Text Text
breakStripPoint :: Text -> Text -> Tuple2 Text Text
breakStripPoint Text
point Text
txt = case HasCallStack => Text -> Text -> Tuple2 Text Text
Text -> Text -> Tuple2 Text Text
T.breakOn Text
point Text
txt of
  (Text
x, Text -> Text -> Maybe Text
T.stripPrefix Text
point -> Just Text
y) -> (Text
x, Text
y)
  Tuple2 Text Text
pair -> Tuple2 Text Text
pair

-- | For 'Natural' \(n\) and 'Text' \(t = t_0 t_1 \ldots t_m\), truncates
-- \(t\) if \(m > n\). In this case, \(t\) is truncated to \(n - 3\), and an
-- ellipsis ( \(\ldots\) ) is appended. We are left with a string with
-- length exactly \(n\):
--
-- \[
-- t_0 t_1 \ldots t_{n-3} \text{...} \quad \text{-- 3 literal } `\text{.' chars appended}
-- \]
--
-- ==== __Examples__
-- >>> truncateIfNeeded 7 "hi"
-- "hi"
--
-- >>> truncateIfNeeded 10 "This is 21 chars long"
-- "This is..."
truncateIfNeeded :: Natural -> Text -> Text
truncateIfNeeded :: Natural -> Text -> Text
truncateIfNeeded Natural
n Text
txt
  | Text -> Int
T.length Text
txt Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
n' = Text
txt
  | Bool
otherwise = Text
txt'
  where
    txt' :: Text
txt' = Int -> Text -> Text
T.take (Int
n' Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
3) Text
txt Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"..."
    n' :: Int
n' = Natural -> Int
n2i Natural
n

n2i :: Natural -> Int
n2i :: Natural -> Int
n2i = Natural -> Int
forall a b.
(Bits a, Bits b, HasCallStack, Integral a, Integral b, Show a,
 Typeable a, Typeable b) =>
a -> b
unsafeConvertIntegral

-- NOTE: [StripControl Newlines]
--
-- Applying stripControl to text that has newlines in it can produce poorly
-- formatted text. This is due to newlines being stripped, so e.g. t1\nt2
-- becomes t1t2. Hence we require 'UnlinedText'.

-- | Strips all control chars, including ansi escape sequences.
--
-- ==== __Examples__
--
-- >>> stripControlAll "foo\ESC[0;3Abar \n baz"
-- "foobar  baz"
stripControlAll :: UnlinedText -> UnlinedText
stripControlAll :: UnlinedText -> UnlinedText
stripControlAll =
  -- The ansi stripping must come first. For example, if we strip control
  -- chars from "\ESC[0;3mfoo" we get "0;3mfoo", and then stripAnsiAll will
  -- no longer recognize this as an ansi sequences - i.e. this will leave
  -- remnants from the ansi sequences.
  --
  -- By performing stripAnsiAll first, we remove entire ansi sequences,
  -- then remove other control chars (e.g. newlines, tabs).
  (Text -> Text) -> UnlinedText -> UnlinedText
ShrunText.reallyUnsafeLiftUnlined ((Char -> Bool) -> Text -> Text
T.filter (Bool -> Bool
not (Bool -> Bool) -> (Char -> Bool) -> Char -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Char -> Bool
isControl) (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
stripAnsiAll)

-- | Strips control chars, including most ansi escape sequences. We leave
-- behind SGR ansi escape sequences e.g. text coloring. See
-- https://en.wikipedia.org/wiki/ANSI_escape_code#SGR_(Select_Graphic_Rendition)_parameters.
--
-- ==== __Examples__
--
-- >>> stripControlSmart "foo\ESC[0;3Abar \n baz"
-- "foobar  baz"
--
-- >>> stripControlSmart "foo\ESC[0;3mbar \n baz"
-- "foo\ESC[0;3mbar  baz"
stripControlSmart :: UnlinedText -> UnlinedText
stripControlSmart :: UnlinedText -> UnlinedText
stripControlSmart =
  -- Like 'stripControlAll', we need to handle the ansi sequences first.
  -- Because we actually leave some sequences behind, we need to be more
  -- surgical removing the rest of the control chars (e.g. newline, tabs).
  (Text -> Text) -> UnlinedText -> UnlinedText
ShrunText.reallyUnsafeLiftUnlined ((Char -> Bool) -> Text -> Text
T.filter Char -> Bool
ctrlToFilter (Text -> Text) -> (Text -> Text) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
stripAnsiControl)
  where
    -- stripAnsiControl should be handling all \ESC sequences, so we should
    -- be safe to ignore these, accomplishing our goal of preserving the SGR
    -- sequences. If this is too aggressive, we can instead attempt to strip
    -- out the known 'bad' control chars e.g.
    --
    --   ctrlToFilter = not . (`elem` ['\n', '\t', '\v'])
    --
    ctrlToFilter :: Char -> Bool
ctrlToFilter Char
c
      | Char -> Bool
isControl Char
c = Char
c Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\ESC'
      | Bool
otherwise = Bool
True

-- | Strips all ansi sequences from the given text.
--
-- ==== __Examples__
--
-- @
-- stripAnsiAll "foo\ESC[0;3Abar"
-- @
stripAnsiAll :: Text -> Text
stripAnsiAll :: Text -> Text
stripAnsiAll = [Text] -> Text
T.concat ([Text] -> Text) -> (Text -> [Text]) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ((Text, Text, Text) -> Text) -> [(Text, Text, Text)] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
fmap (Optic' A_Lens NoIx (Text, Text, Text) Text
-> (Text, Text, Text) -> Text
forall k (is :: IxList) s a.
Is k A_Getter =>
Optic' k is s a -> s -> a
view Optic' A_Lens NoIx (Text, Text, Text) Text
forall s t a b. Field1 s t a b => Lens s t a b
_1) ([(Text, Text, Text)] -> [Text])
-> (Text -> [(Text, Text, Text)]) -> Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> [(Text, Text, Text)]
splitAnsi

-- | Strips ansi control sequences only.
--
-- ==== __Examples__
--
-- @
-- stripAnsiControl "foo\ESC[0;3Abar"
-- "foobar"
--
-- stripAnsiControl "foo\ESC[0;3mbar"
-- "foo\ESC[0;3mbar"
-- @
stripAnsiControl :: Text -> Text
stripAnsiControl :: Text -> Text
stripAnsiControl Text
txt =
  (Text -> (Text, Text, Text) -> Text)
-> Text -> [(Text, Text, Text)] -> Text
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: Type -> Type) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' Text -> (Text, Text, Text) -> Text
forall {a}. Semigroup a => a -> (a, Text, a) -> a
f Text
"" [(Text, Text, Text)]
splitTxt
  where
    splitTxt :: [(Text, Text, Text)]
splitTxt = Text -> [(Text, Text, Text)]
splitAnsi Text
txt
    f :: a -> (a, Text, a) -> a
f a
acc (a
preAnsi, Text
code, a
withAnsi)
      | Text -> Bool
nonControlAnsi Text
code = a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
withAnsi
      | Bool
otherwise = a
acc a -> a -> a
forall a. Semigroup a => a -> a -> a
<> a
preAnsi

nonControlAnsi :: Text -> Bool
nonControlAnsi :: Text -> Bool
nonControlAnsi Text
ansi = case Text -> Maybe (Text, Char)
T.unsnoc Text
ansi of
  -- 'm' equals color: only code we consider 'good' for now
  Just (Text
_, Char
'm') -> Bool
True
  Maybe (Text, Char)
_ -> Bool
False

-- tuple is: (text, ansi_code, ansi_code <> text)
-- example: splitAnsi "foo\ESC[0;3mbar"
splitAnsi :: Text -> [(Text, Text, Text)]
splitAnsi :: Text -> [(Text, Text, Text)]
splitAnsi Text
"" = []
splitAnsi Text
t =
  -- (foo, \ESC[0;3m, bar) : ...
  (Text
preAnsi, Text
ansiCode, Text
preAnsi Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
ansiCode) (Text, Text, Text) -> [(Text, Text, Text)] -> [(Text, Text, Text)]
forall a. a -> [a] -> [a]
: [(Text, Text, Text)]
rest
  where
    -- (foo, \ESC[0;3mbar)
    (!Text
preAnsi, !Text
withAnsiFull) = HasCallStack => Text -> Text -> Tuple2 Text Text
Text -> Text -> Tuple2 Text Text
T.breakOn Text
"\ESC[" Text
t
    -- (\ESC[0;3, mbar)
    (!Text
ansiCodeNoChar, !Text
withAnsiChar) = (Char -> Bool) -> Text -> Tuple2 Text Text
T.break Char -> Bool
isLetter Text
withAnsiFull
    -- (\ESC[0;3m, bar)
    (!Text
ansiCode, ![(Text, Text, Text)]
rest) = case Text -> Maybe (Char, Text)
T.uncons Text
withAnsiChar of
      -- (m, bar)
      Just (!Char
ansiChar, !Text
rest') -> (Text -> Char -> Text
T.snoc Text
ansiCodeNoChar Char
ansiChar, Text -> [(Text, Text, Text)]
splitAnsi Text
rest')
      Maybe (Char, Text)
Nothing -> (Text
ansiCodeNoChar, [])

-- | Parses bytes with arbitrary units and converts to bytes. First attempts
-- to parse as a 'Natural' so we do not lose precision. If that fails, falls
-- back to 'Double'.
--
-- ==== __Examples__
--
-- >>> parseByteText "120 mb"
-- Right (MkBytes 120000000)
--
-- >>> parseByteText "4.5 terabytes"
-- Right (MkBytes 4500000000000)
parseByteText :: Text -> Either Text (Bytes B Natural)
parseByteText :: Text -> Either Text (Bytes 'B Natural)
parseByteText Text
txt =
  case forall a. Parser a => Text -> Either Text a
parse @(SomeSize Natural) Text
txt of
    Right SomeSize Natural
b -> Bytes 'B Natural -> Either Text (Bytes 'B Natural)
forall a b. b -> Either a b
Right (Bytes 'B Natural -> Either Text (Bytes 'B Natural))
-> Bytes 'B Natural -> Either Text (Bytes 'B Natural)
forall a b. (a -> b) -> a -> b
$ Proxy 'B -> SomeSize Natural -> Converted 'B (SomeSize Natural)
forall a (t :: Size).
(Conversion a, SingSize t) =>
Proxy t -> a -> Converted t a
forall (t :: Size).
SingSize t =>
Proxy t -> SomeSize Natural -> Converted t (SomeSize Natural)
convert (forall {k} (t :: k). Proxy t
forall (t :: Size). Proxy t
Proxy @B) SomeSize Natural
b
    Left Text
_ -> case forall a. Parser a => Text -> Either Text a
parse @(SomeSize Double) Text
txt of
      Right SomeSize Double
b -> Bytes 'B Natural -> Either Text (Bytes 'B Natural)
forall a b. b -> Either a b
Right (Double -> Natural
forall b. Integral b => Double -> b
forall a b. (RealFrac a, Integral b) => a -> b
truncate (Double -> Natural) -> Bytes 'B Double -> Bytes 'B Natural
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Proxy 'B -> SomeSize Double -> Converted 'B (SomeSize Double)
forall a (t :: Size).
(Conversion a, SingSize t) =>
Proxy t -> a -> Converted t a
forall (t :: Size).
SingSize t =>
Proxy t -> SomeSize Double -> Converted t (SomeSize Double)
convert (forall {k} (t :: k). Proxy t
forall (t :: Size). Proxy t
Proxy @B) SomeSize Double
b)
      Left Text
err -> Text -> Either Text (Bytes 'B Natural)
forall a b. a -> Either a b
Left Text
err

-- | Runs the action when it is 'Left'.
whenLeft :: (Applicative f) => Either a b -> (a -> f ()) -> f ()
whenLeft :: forall (f :: Type -> Type) a b.
Applicative f =>
Either a b -> (a -> f ()) -> f ()
whenLeft Either a b
e a -> f ()
action = (a -> f ()) -> (b -> f ()) -> Either a b -> f ()
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either a -> f ()
action (f () -> b -> f ()
forall a b. a -> b -> a
const (() -> f ()
forall a. a -> f a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ())) Either a b
e

-- | @whileM_ mb ma@ executes @ma@ as long as @mb@ returns 'True'.
whileM_ :: (Monad m) => m Bool -> m a -> m ()
whileM_ :: forall (m :: Type -> Type) a. Monad m => m Bool -> m a -> m ()
whileM_ m Bool
mb m a
ma = m ()
go
  where
    go :: m ()
go =
      m Bool
mb m Bool -> (Bool -> m ()) -> m ()
forall a b. m a -> (a -> m b) -> m b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Bool
True -> m a
ma m a -> m () -> m ()
forall a b. m a -> m b -> m b
forall (f :: Type -> Type) a b. Applicative f => f a -> f b -> f b
*> m ()
go
        Bool
False -> () -> m ()
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure ()

-- | Executes the monadic action until we receive a 'Just', returning the
-- value.
untilJust :: (Monad m) => m (Maybe b) -> m b
untilJust :: forall (m :: Type -> Type) b. Monad m => m (Maybe b) -> m b
untilJust m (Maybe b)
m = m b
go
  where
    go :: m b
go =
      m (Maybe b)
m m (Maybe b) -> (Maybe b -> m b) -> m b
forall a b. m a -> (a -> m b) -> m b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
        Maybe b
Nothing -> m b
go
        Just b
x -> b -> m b
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure b
x

{- HLINT ignore unsafeListToNESeq "Redundant bracket" -}

unsafeListToNESeq :: (HasCallStack) => List a -> NESeq a
unsafeListToNESeq :: forall a. HasCallStack => List a -> NESeq a
unsafeListToNESeq [] = [Char] -> NESeq a
forall a. HasCallStack => [Char] -> a
error [Char]
"[Shrun.Utils]: empty list"
unsafeListToNESeq [a]
xs = NonEmpty a -> NESeq a
forall a. NonEmpty a -> NESeq a
NESeq.fromList (NonEmpty a -> NESeq a) -> NonEmpty a -> NESeq a
forall a b. (a -> b) -> a -> b
$ [Item (NonEmpty a)] -> NonEmpty a
forall l. IsList l => [Item l] -> l
fromList [a]
[Item (NonEmpty a)]
xs

-- | Escape double quotes in strings.
escapeDoubleQuotes :: Text -> Text
escapeDoubleQuotes :: Text -> Text
escapeDoubleQuotes = LazyText -> Text
TL.toStrict (LazyText -> Text) -> (Text -> LazyText) -> Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Builder -> LazyText
TLB.toLazyText (Builder -> LazyText) -> (Text -> Builder) -> Text -> LazyText
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Builder -> Char -> Builder) -> Builder -> Text -> Builder
forall a. (a -> Char -> a) -> a -> Text -> a
T.foldl' Builder -> Char -> Builder
go Builder
""
  where
    go :: Builder -> Char -> Builder
    go :: Builder -> Char -> Builder
go Builder
acc Char
'"' = Builder
acc Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Builder
"\\\""
    go Builder
acc Char
c = Builder
acc Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> Char -> Builder
TLB.singleton Char
c

-- | "monus" i.e. subtraction clamped to zero
(∸) :: (Ord a, Num a) => a -> a -> a
a
x ∸ :: forall a. (Ord a, Num a) => a -> a -> a
∸ a
y =
  if a
y a -> a -> Bool
forall a. Ord a => a -> a -> Bool
> a
x
    then a
0
    else a
x a -> a -> a
forall a. Num a => a -> a -> a
- a
y

infixl 6 ∸

readStripUnderscores :: (MonadFail m, Read a) => Text -> m a
readStripUnderscores :: forall (m :: Type -> Type) a. (MonadFail m, Read a) => Text -> m a
readStripUnderscores Text
t = case [Char] -> Either [Char] a
forall a. Read a => [Char] -> Either [Char] a
TR.readEither [Char]
s of
  Left [Char]
err -> [Char] -> m a
forall a. [Char] -> m a
forall (m :: Type -> Type) a. MonadFail m => [Char] -> m a
fail ([Char] -> m a) -> [Char] -> m a
forall a b. (a -> b) -> a -> b
$ [Char]
"Could not read '" [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
s [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
"': " [Char] -> [Char] -> [Char]
forall a. [a] -> [a] -> [a]
++ [Char]
err
  Right a
x -> a -> m a
forall a. a -> m a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure a
x
  where
    noUnderscores :: Text
noUnderscores = HasCallStack => Text -> Text -> Text -> Text
Text -> Text -> Text -> Text
T.replace Text
"_" Text
"" Text
t
    s :: [Char]
s = Text -> [Char]
T.unpack Text
noUnderscores