module Distribution.Nixpkgs.Color
    ( maybeColor
    , colorStderrLn
    , infoColor
    , warningColor
    , errorColor
    , commandColor
    ) where

import System.Environment (lookupEnv)
import System.IO (Handle, hIsTerminalDevice, hPutStrLn, stderr)
import System.Console.ANSI.Codes
    ( setSGRCode
    , SGR(Reset, SetColor, SetConsoleIntensity)
    , ConsoleLayer(Foreground)
    , ColorIntensity(Vivid)
    , Color(Yellow, Red, Cyan)
    , ConsoleIntensity(BoldIntensity)
    )
import Control.Monad.IO.Class (MonadIO(liftIO))

-- | Colors that indicate a warning message.
warningColor :: [SGR]
warningColor :: [SGR]
warningColor = [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Vivid Color
Yellow]

infoColor :: [SGR]
infoColor :: [SGR]
infoColor = [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Vivid Color
Cyan]

-- | Colors that indicate an error message.
errorColor :: [SGR]
errorColor :: [SGR]
errorColor = [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Vivid Color
Red, ConsoleIntensity -> SGR
SetConsoleIntensity ConsoleIntensity
BoldIntensity]

-- | Colors that indicate a command is being executed.
commandColor :: [SGR]
commandColor :: [SGR]
commandColor = [ConsoleLayer -> ColorIntensity -> Color -> SGR
SetColor ConsoleLayer
Foreground ColorIntensity
Vivid Color
Cyan, ConsoleIntensity -> SGR
SetConsoleIntensity ConsoleIntensity
BoldIntensity]

-- | Check if an environment variable is set and non-empty.
envIsSet :: String -> IO Bool
envIsSet :: String -> IO Bool
envIsSet String
name = do
    Maybe String
value <- String -> IO (Maybe String)
lookupEnv String
name
    Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Bool -> IO Bool) -> Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ case Maybe String
value of
       Maybe String
Nothing -> Bool
False
       Just String
"" -> Bool
False
       Just String
_ -> Bool
True

-- | Should output to the given `Handle` be colored?
shouldColor :: Handle -> IO Bool
shouldColor :: Handle -> IO Bool
shouldColor Handle
handle = do
    -- See: https://no-color.org/
    Bool
noColor <- String -> IO Bool
envIsSet String
"NO_COLOR"
    if Bool
noColor
       then Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
False
       else do
           Bool
forceColor <- String -> IO Bool
envIsSet String
"FORCE_COLOR"
           if Bool
forceColor
              then Bool -> IO Bool
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True
              else Handle -> IO Bool
hIsTerminalDevice Handle
handle

-- | If the given `Handle` should be colored, wrap a `String` in `SGR` codes.
maybeColor :: Handle -> [SGR] -> String -> IO String
maybeColor :: Handle -> [SGR] -> String -> IO String
maybeColor Handle
handle [SGR]
sgrCodes String
original = do
    Bool
shouldColor' <- Handle -> IO Bool
shouldColor Handle
handle
    if Bool -> Bool
not Bool
shouldColor'
       then String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure String
original
       else String -> IO String
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> IO String) -> String -> IO String
forall a b. (a -> b) -> a -> b
$ [SGR] -> String
setSGRCode [SGR]
sgrCodes String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
original String -> String -> String
forall a. Semigroup a => a -> a -> a
<> [SGR] -> String
setSGRCode [SGR
Reset]

colorStderrLn :: MonadIO m => [SGR] -> String -> m ()
colorStderrLn :: forall (m :: * -> *). MonadIO m => [SGR] -> String -> m ()
colorStderrLn [SGR]
sgrCodes String
original = IO () -> m ()
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$ do
    String
maybeColored <- Handle -> [SGR] -> String -> IO String
maybeColor Handle
stderr [SGR]
sgrCodes String
original
    Handle -> String -> IO ()
hPutStrLn Handle
stderr String
maybeColored