module Rainbow.Internal where

import Control.Exception (try)
import Data.ByteString.Builder (Builder, char7, word8Dec, byteString, hPutBuilder)
import Data.List (intersperse)
import Data.String (IsString(..))
import Data.Text (Text, pack)
import Data.Text.Encoding (encodeUtf8)
import Data.Word (Word8)
import qualified System.Console.Terminfo as Terminfo
import qualified System.IO as IO

--
-- Chunks
--

-- | A chunk is some textual data coupled with a description of what
-- color the text is, attributes like whether it is bold or
-- underlined, etc. The chunk knows what foreground and background
-- colors and what attributes to use.

data Chunk = Chunk
  { format :: Builder
  , content :: Text
  -- ^ Returns the textual content of the 'Chunk'.
  }

-- | The 'Show' instance shows only the contained 'Text' and none of the formatting.
instance Show Chunk where
  show (Chunk _ txt) = show txt

-- | Creates a 'Chunk' with no formatting and with the given text.
instance IsString Chunk where
  fromString = Chunk mempty . pack

-- | Creates a 'Chunk' with no formatting and with the given text.
-- A 'Chunk' is also an instance of 'Data.String.IsString' so you
-- can create them with the @OverloadedStrings@ extension.  Such a
-- 'Chunk' has the text of the string and no formatting.
chunk :: Text -> Chunk
chunk = Chunk mempty

applyFormat :: Builder -> Chunk -> Chunk
applyFormat b (Chunk f t) = Chunk (f <> b) t

-- | Bold. What actually happens when you use Bold is going to depend
-- on your terminal. For example, xterm allows you actually use a bold
-- font for bold, if you have one. Otherwise, it might simulate bold
-- by using overstriking. Another possibility is that your terminal
-- might use a different color to indicate bold. For more details (at
-- least for xterm), look at xterm (1) and search for @boldColors@.
--
-- If your terminal uses a different color for bold, this allows an
-- 8-color terminal to really have 16 colors.
bold :: Chunk -> Chunk
bold = applyFormat boldF

faint :: Chunk -> Chunk
faint = applyFormat faintF

italic :: Chunk -> Chunk
italic = applyFormat italicF

underline :: Chunk -> Chunk
underline = applyFormat underlineF

blink :: Chunk -> Chunk
blink = applyFormat blinkF

inverse :: Chunk -> Chunk
inverse = applyFormat inverseF

invisible :: Chunk -> Chunk
invisible = applyFormat invisibleF

strikeout :: Chunk -> Chunk
strikeout = applyFormat strikeoutF

data Pigment
  = Black
  | Red
  | Green
  | Yellow
  | Blue
  | Magenta
  | Cyan
  | White
  deriving (Eq, Ord, Show)


-- | One of 8 colors, and whether or not the color is the bright variant.
data Color3bit = Color3bit
  { pigment :: Pigment
  , isBright :: Bool
  } deriving (Eq, Ord, Show)

-- | A color for display on a 256-color terminal.
data Color
  = Color8 Color3bit
  | Color256 Word8
  deriving (Eq, Ord, Show)

black, red, green, yellow, blue, magenta, cyan, white :: Color
black = Color8 $ Color3bit Black False
red = Color8 $ Color3bit Red False
green = Color8 $ Color3bit Green False
yellow = Color8 $ Color3bit Yellow False
blue = Color8 $ Color3bit Blue False
magenta = Color8 $ Color3bit Magenta False
cyan = Color8 $ Color3bit Cyan False
white = Color8 $ Color3bit White False

-- | Use the bright variant of a color.  Has an effect only on colors
-- constructed with 'black', 'red', 'green', 'yellow', 'blue', 'magenta',
-- 'cyan', or 'white'; has no effect on colors constructed with 'color256'.
bright :: Color -> Color
bright (Color8 (Color3bit p _)) = Color8 (Color3bit p True)
bright x = x

-- | A 'Color' for any of the 256 colors available.  Supply the
-- color number.  Exactly which color you'll get for a given number
-- is dependent on the terminal; though there seem to be common
-- defaults, often the user can configure this however she likes.
color256 :: Word8 -> Color
color256 = Color256

-- | Change the foreground color.

fore :: Color -> Chunk -> Chunk
fore clr (Chunk f t) = Chunk (f <> c) t
  where
    c = case clr of
      Color8 (Color3bit pig isBri)
        | isBri -> case pig of
            Black -> foreBrightBlack
            Red -> foreBrightRed
            Green -> foreBrightGreen
            Yellow -> foreBrightYellow
            Blue -> foreBrightBlue
            Magenta -> foreBrightMagenta
            Cyan -> foreBrightCyan
            White -> foreBrightWhite
        | otherwise -> case pig of
            Black -> foreBlack
            Red -> foreRed
            Green -> foreGreen
            Yellow -> foreYellow
            Blue -> foreBlue
            Magenta -> foreMagenta
            Cyan -> foreCyan
            White -> foreWhite
      Color256 c256 -> fore256 c256

-- | Change the background color.
back :: Color -> Chunk -> Chunk
back clr (Chunk f t) = Chunk (f <> c) t
  where
    c = case clr of
      Color8 (Color3bit pig isBri)
        | isBri -> case pig of
            Black -> backBrightBlack
            Red -> backBrightRed
            Green -> backBrightGreen
            Yellow -> backBrightYellow
            Blue -> backBrightBlue
            Magenta -> backBrightMagenta
            Cyan -> backBrightCyan
            White -> backBrightWhite
        | otherwise -> case pig of
            Black -> backBlack
            Red -> backRed
            Green -> backGreen
            Yellow -> backYellow
            Blue -> backBlue
            Magenta -> backMagenta
            Cyan -> backCyan
            White -> backWhite
      Color256 c256 -> back256 c256

-- | Same as @'bright' 'black'@.
grey :: Color
grey = bright black

single :: Char -> Builder
single = char7

escape :: Builder
escape = single '\x1B'

csi :: Builder
csi = escape <> single '['

sgr :: Builder -> Builder
sgr sq = csi <> sq <> single 'm'

params :: [Word8] -> Builder
params = mconcat . intersperse (char7 ';') . map word8Dec

sgrSingle :: Word8 -> Builder
sgrSingle w = sgr $ params [w]

sgrDouble :: Word8 -> Word8 -> Builder
sgrDouble x y = sgr $ params [x, y]

normalDefault :: Builder
normalDefault = sgrSingle 0

boldF :: Builder
boldF = sgrSingle 1

faintF :: Builder
faintF = sgrSingle 2

italicF :: Builder
italicF = sgrSingle 3

underlineF :: Builder
underlineF = sgrSingle 4

blinkF :: Builder
blinkF = sgrSingle 5

-- Yes, blink is 5, inverse is 7; 6 is skipped.  In ISO 6429 6 blinks
-- at a different rate.

inverseF :: Builder
inverseF = sgrSingle 7

invisibleF :: Builder
invisibleF = sgrSingle 8

strikeoutF :: Builder
strikeoutF = sgrSingle 9

foreBlack :: Builder
foreBlack = sgrSingle 30

foreRed :: Builder
foreRed = sgrSingle 31

foreGreen :: Builder
foreGreen = sgrSingle 32

foreYellow :: Builder
foreYellow = sgrSingle 33

foreBlue :: Builder
foreBlue = sgrSingle 34

foreMagenta :: Builder
foreMagenta = sgrSingle 35

foreCyan :: Builder
foreCyan = sgrSingle 36

foreWhite :: Builder
foreWhite = sgrSingle 37

foreBrightBlack :: Builder
foreBrightBlack = sgrSingle 00

foreBrightRed :: Builder
foreBrightRed = sgrSingle 91

foreBrightGreen :: Builder
foreBrightGreen = sgrSingle 92

foreBrightYellow :: Builder
foreBrightYellow = sgrSingle 93

foreBrightBlue :: Builder
foreBrightBlue = sgrSingle 94

foreBrightMagenta :: Builder
foreBrightMagenta = sgrSingle 95

foreBrightCyan :: Builder
foreBrightCyan = sgrSingle 96

foreBrightWhite :: Builder
foreBrightWhite = sgrSingle 97

-- code 3 8 is skipped

foreDefault :: Builder
foreDefault = sgrSingle 39

backBlack :: Builder
backBlack = sgrSingle 40

backRed :: Builder
backRed = sgrSingle 41

backGreen :: Builder
backGreen = sgrSingle 42

backYellow :: Builder
backYellow = sgrSingle 43

backBlue :: Builder
backBlue = sgrSingle 44

backMagenta :: Builder
backMagenta = sgrSingle 45

backCyan :: Builder
backCyan = sgrSingle 46

backWhite :: Builder
backWhite = sgrSingle 47

backBrightBlack :: Builder
backBrightBlack = sgrSingle 100

backBrightRed :: Builder
backBrightRed = sgrSingle 101

backBrightGreen :: Builder
backBrightGreen = sgrSingle 102

backBrightYellow :: Builder
backBrightYellow = sgrSingle 103

backBrightBlue :: Builder
backBrightBlue = sgrSingle 104

backBrightMagenta :: Builder
backBrightMagenta = sgrSingle 105

backBrightCyan :: Builder
backBrightCyan = sgrSingle 106

backBrightWhite :: Builder
backBrightWhite = sgrSingle 107

-- code 4 8 is skipped

backDefault :: Builder
backDefault = sgrSingle 49

fore256 :: Word8 -> Builder
fore256 c = sgr $ params [38,5,c]

back256 :: Word8 -> Builder
back256 c = sgr $ params [48,5,c]


render :: Text -> Builder
render = byteString . encodeUtf8

toBuilderColors0 :: Chunk -> Builder
toBuilderColors0 (Chunk _ yn) = render yn

toBuilderColors256 :: Chunk -> Builder
toBuilderColors256 (Chunk fmt yn)
  = normalDefault <> fmt <> render yn <> normalDefault


-- | Uses 'Terminfo.setupTermFromEnv' to obtain the terminal's color
-- capability.  If this says there are at least 256 colors are
-- available, returns 'toBuilderColors256'.
-- Otherwise, returns 'toBuilderColors0'.
--
-- If the terminfo database could not be read (that is, if
-- 'System.Console.Terminfo.Base.SetupTermError' is returned), then return
-- 'toBuilderColors0'.
builderFromEnvironment :: IO (Chunk -> Builder)
builderFromEnvironment = fmap g (try Terminfo.setupTermFromEnv)
  where
    g (Left e) = toBuilderColors0
      where
        -- Previously this caught all IOException.  Now it catches only SetupTermError.
        -- See
        -- https://github.com/commercialhaskell/stackage/issues/4994
        -- Hopefully this will fix this Stackage bug.
        _types = e :: Terminfo.SetupTermError
    g (Right terminal) =
      case Terminfo.getCapability terminal (Terminfo.tiGetNum "colors") of
        Nothing -> toBuilderColors0
        Just c
          | c >= 256 -> toBuilderColors256
          | otherwise -> toBuilderColors0

-- | Like 'builderFromEnvironment' but also consults a
-- provided 'IO.Handle'.  If the 'IO.Handle' is not a terminal,
-- 'toBuilderColors0' is returned.  Otherwise, the value of
-- 'builderFromEnvironment' is returned.
builderFromHandle
  :: IO.Handle
  -> IO (Chunk -> Builder)
builderFromHandle h = IO.hIsTerminalDevice h >>= f
  where
    f isTerm | isTerm = builderFromEnvironment
             | otherwise = return toBuilderColors0
          

-- | Convert a list of 'Chunk' to a 'Builder'.

chunksToBuilder
  :: (Chunk -> Builder)
  -- ^ Function that converts 'Chunk' to 'Builder'.
  -> [Chunk]
  -> Builder
chunksToBuilder mk = mconcat . map mk

-- | Writes a list of chunks to the given 'IO.Handle'.
--
-- First uses 'builderFromEnvironment' to determine how many colors to
-- use.  Sets the 'IO.Handle' binary and BlockBuffering mode.  Then creates a
-- Builder using 'chunksToBuilder' and then writes it to the given
-- 'IO.Handle'.
hPutChunks :: IO.Handle -> [Chunk] -> IO ()
hPutChunks h cks = do
  maker <- builderFromEnvironment
  IO.hSetBinaryMode h True
  IO.hSetBuffering h (IO.BlockBuffering Nothing)
  hPutBuilder h . chunksToBuilder maker $ cks

-- | Writes a list of chunks to the given 'IO.Handle', followed by a
-- newline character.
--
-- First uses 'builderFromEnvironment' to determine how many
-- colors to use.  Then creates a list of 'ByteString' using
-- 'chunksToBuilder' and then writes them to the given 'IO.Handle'.
hPutChunksLn :: IO.Handle -> [Chunk] -> IO ()
hPutChunksLn h cks = do
  hPutChunks h cks
  IO.hPutStr h "\n"

-- | Writes a list of chunks to standard output.
--
-- First uses 'builderFromEnvironment' to determine how many
-- colors to use.  Then creates a list of 'ByteString' using
-- 'chunksToBuilder' and then writes them to standard output.
putChunks :: [Chunk] -> IO ()
putChunks = hPutChunks IO.stdout

-- | Writes a list of chunks to standard output, followed by a
-- newline.
--
-- First uses 'builderFromEnvironment' to determine how many
-- colors to use.  Then creates a list of 'ByteString' using
-- 'chunksToBuilder' and then writes them to standard output.
putChunksLn :: [Chunk] -> IO ()
putChunksLn cks = do
  putChunks cks
  IO.putStr "\n"

-- | Writes a 'Chunk' to standard output.  Uses
-- 'builderFromEnvironment' each time you apply it, so this
-- might be inefficient.  You are better off using
-- 'chunksToBuilder' and the functions in "Data.ByteString" to
-- print your 'Chunk's if you are printing a lot of them.
putChunk :: Chunk -> IO ()
putChunk ck = hPutChunks IO.stdout [ck]

-- | Writes a 'Chunk' to standard output, and appends a newline.
-- Uses 'builderFromEnvironment' each time you apply it, so
-- this might be inefficient.  You are better off using
-- 'chunksToBuilder' and the functions in "Data.ByteString" to
-- print your 'Chunk's if you are printing a lot of them.
putChunkLn :: Chunk -> IO ()
putChunkLn ck = putChunk ck >> putStrLn ""