{-# LANGUAGE OverloadedStrings #-}
module Vgrep.Ansi.Parser
  ( parseAnsi
  , ansiFormatted
  , attrChange
  ) where


import           Control.Applicative
import           Data.Attoparsec.Text
import           Data.Bits
import           Data.Functor
import           Data.Text               (Text)
import qualified Data.Text               as T
import           Graphics.Vty.Attributes (Attr)
import qualified Graphics.Vty.Attributes as Vty

import Vgrep.Ansi.Type


{- |
Directly parses ANSI formatted text using 'ansiFormatted'.

Parsing ANSI color codes:

>>> parseAnsi "Hello \ESC[31mWorld\ESC[m!"
Cat 12 [Text 6 "Hello ",Format 5 (Attr {attrStyle = KeepCurrent, attrForeColor = SetTo (ISOColor 1), attrBackColor = KeepCurrent, attrURL = KeepCurrent}) (Text 5 "World"),Text 1 "!"]

More elaborate example with nested foreground and background colors:

>>> parseAnsi "\ESC[m\ESC[40mHello \ESC[31mWorld\ESC[39m!"
Cat 12 [Format 6 (Attr {attrStyle = KeepCurrent, attrForeColor = KeepCurrent, attrBackColor = SetTo (ISOColor 0), attrURL = KeepCurrent}) (Text 6 "Hello "),Format 5 (Attr {attrStyle = KeepCurrent, attrForeColor = SetTo (ISOColor 1), attrBackColor = SetTo (ISOColor 0), attrURL = KeepCurrent}) (Text 5 "World"),Format 1 (Attr {attrStyle = KeepCurrent, attrForeColor = KeepCurrent, attrBackColor = SetTo (ISOColor 0), attrURL = KeepCurrent}) (Text 1 "!")]

Some CSI sequences are ignored, since they are not supported by 'Vty':

>>> parseAnsi "\ESC[A\ESC[B\ESC[31mfoo\ESC[1K\ESC[mbar"
Cat 6 [Format 3 (Attr {attrStyle = KeepCurrent, attrForeColor = SetTo (ISOColor 1), attrBackColor = KeepCurrent, attrURL = KeepCurrent}) (Text 3 "foo"),Text 3 "bar"]

Non-CSI sequences are not parsed, but included in the output:

>>> parseAnsi "\ESC]710;font\007foo\ESC[31mbar"
Cat 17 [Text 14 "\ESC]710;font\afoo",Format 3 (Attr {attrStyle = KeepCurrent, attrForeColor = SetTo (ISOColor 1), attrBackColor = KeepCurrent, attrURL = KeepCurrent}) (Text 3 "bar")]

-}
parseAnsi :: Text -> AnsiFormatted
parseAnsi :: Text -> AnsiFormatted
parseAnsi = ([Char] -> AnsiFormatted)
-> (AnsiFormatted -> AnsiFormatted)
-> Either [Char] AnsiFormatted
-> AnsiFormatted
forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either [Char] -> AnsiFormatted
forall a. HasCallStack => [Char] -> a
error AnsiFormatted -> AnsiFormatted
forall a. a -> a
id (Either [Char] AnsiFormatted -> AnsiFormatted)
-> (Text -> Either [Char] AnsiFormatted) -> Text -> AnsiFormatted
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Parser AnsiFormatted -> Text -> Either [Char] AnsiFormatted
forall a. Parser a -> Text -> Either [Char] a
parseOnly Parser AnsiFormatted
ansiFormatted
-- The use of 'error' ↑ is safe: 'ansiFormatted' does not fail.


-- | Parser for ANSI formatted text. Recognized escape sequences are the SGR
-- (Select Graphic Rendition) sequences (@\ESC[…m@) supported by 'Attr'.
-- Unsupported SGR sequences and other CSI escape sequences (@\ESC[…@) are
-- ignored. Other (non-CSI) escape sequences are not parsed, and included in the
-- output.
--
-- This parser does not fail, it will rather consume and return the remaining
-- input as unformatted text.
ansiFormatted :: Parser AnsiFormatted
ansiFormatted :: Parser AnsiFormatted
ansiFormatted = Attr -> Parser AnsiFormatted
go Attr
forall a. Monoid a => a
mempty
  where
    go :: Attr -> Parser AnsiFormatted
    go :: Attr -> Parser AnsiFormatted
go Attr
attr = Parser Text ()
forall t. Chunk t => Parser t ()
endOfInput Parser Text () -> AnsiFormatted -> Parser AnsiFormatted
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> AnsiFormatted
forall a. Monoid a => a
mempty
          Parser AnsiFormatted
-> Parser AnsiFormatted -> Parser AnsiFormatted
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Attr -> Parser AnsiFormatted
formattedText Attr
attr

    formattedText :: Attr -> Parser AnsiFormatted
    formattedText :: Attr -> Parser AnsiFormatted
formattedText Attr
attr = do
        [Attr -> Attr]
acs <- Parser Text (Attr -> Attr) -> Parser Text [Attr -> Attr]
forall (f :: * -> *) a. Alternative f => f a -> f [a]
many Parser Text (Attr -> Attr)
attrChange
        let attr' :: Attr
attr' = ((Attr -> Attr) -> Attr -> Attr) -> Attr -> [Attr -> Attr] -> Attr
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (Attr -> Attr) -> Attr -> Attr
forall a b. (a -> b) -> a -> b
($) Attr
attr ([Attr -> Attr] -> [Attr -> Attr]
forall a. [a] -> [a]
reverse [Attr -> Attr]
acs)
        Text
t <- Parser Text
rawText
        AnsiFormatted
rest <- Attr -> Parser AnsiFormatted
go Attr
attr'
        AnsiFormatted -> Parser AnsiFormatted
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Attr -> AnsiFormatted -> AnsiFormatted
forall attr.
(Eq attr, Monoid attr) =>
attr -> Formatted attr -> Formatted attr
format Attr
attr' (Text -> AnsiFormatted
forall attr. Text -> Formatted attr
bare Text
t) AnsiFormatted -> AnsiFormatted -> AnsiFormatted
forall a. Semigroup a => a -> a -> a
<> AnsiFormatted
rest)

    rawText :: Parser Text
    rawText :: Parser Text
rawText = (Char -> Bool) -> Parser Text
atLeastOneTill (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\ESC') Parser Text -> Parser Text -> Parser Text
forall (f :: * -> *) a. Alternative f => f a -> f a -> f a
<|> Parser Text ()
forall t. Chunk t => Parser t ()
endOfInput Parser Text () -> Text -> Parser Text
forall (f :: * -> *) a b. Functor f => f a -> b -> f b
$> Text
""

    atLeastOneTill :: (Char -> Bool) -> Parser Text
    atLeastOneTill :: (Char -> Bool) -> Parser Text
atLeastOneTill = (Char -> Text -> Text)
-> Parser Text Char -> Parser Text -> Parser Text
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 Char -> Text -> Text
T.cons Parser Text Char
anyChar (Parser Text -> Parser Text)
-> ((Char -> Bool) -> Parser Text) -> (Char -> Bool) -> Parser Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Char -> Bool) -> Parser Text
takeTill


-- | Parser for ANSI CSI escape sequences. Recognized escape sequences are the
-- SGR (Select Graphic Rendition) sequences (@\ESC[…m@) supported by 'Attr'.
-- Unsupported SGR sequences and other CSI escape sequences (@\ESC[…@) are
-- ignored by returning 'id'.
--
-- This parser fails when encountering any other (non-CSI) escape sequence.
attrChange :: Parser (Attr -> Attr)
attrChange :: Parser Text (Attr -> Attr)
attrChange = (Csi -> Attr -> Attr)
-> Parser Text Csi -> Parser Text (Attr -> Attr)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Csi -> Attr -> Attr
csiToAttrChange Parser Text Csi
csi

csiEscape :: Parser Text
csiEscape :: Parser Text
csiEscape = Parser Text
"\ESC["

csi :: Parser Csi
csi :: Parser Text Csi
csi = Parser Text
csiEscape Parser Text -> Parser Text Csi -> Parser Text Csi
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> ([Int] -> Char -> Csi)
-> Parser Text [Int] -> Parser Text Char -> Parser Text Csi
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2 [Int] -> Char -> Csi
Csi (Parser Int
forall a. Integral a => Parser a
decimal Parser Int -> Parser Text Char -> Parser Text [Int]
forall (f :: * -> *) a s. Alternative f => f a -> f s -> f [a]
`sepBy` Char -> Parser Text Char
char Char
';') Parser Text Char
anyChar

data Csi = Csi [Int] Char

csiToAttrChange :: Csi -> Attr -> Attr
csiToAttrChange :: Csi -> Attr -> Attr
csiToAttrChange = \case
    Csi [] Char
'm' -> Attr -> Attr -> Attr
forall a b. a -> b -> a
const Attr
forall a. Monoid a => a
mempty
    Csi [Int]
is Char
'm' -> (Int -> Attr -> Attr) -> [Int] -> Attr -> Attr
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap Int -> Attr -> Attr
attrChangeFromCode [Int]
is
    Csi
_otherwise -> Attr -> Attr
forall a. a -> a
id

attrChangeFromCode :: Int -> Attr -> Attr
attrChangeFromCode :: Int -> Attr -> Attr
attrChangeFromCode = \case
    Int
0  -> Attr -> Attr -> Attr
forall a b. a -> b -> a
const Attr
forall a. Monoid a => a
mempty
    Int
1  -> Style -> Attr -> Attr
withStyle Style
Vty.bold
    Int
3  -> Style -> Attr -> Attr
withStyle Style
Vty.standout
    Int
4  -> Style -> Attr -> Attr
withStyle Style
Vty.underline
    Int
5  -> Style -> Attr -> Attr
withStyle Style
Vty.blink
    Int
6  -> Style -> Attr -> Attr
withStyle Style
Vty.blink
    Int
7  -> Style -> Attr -> Attr
withStyle Style
Vty.reverseVideo
    Int
21 -> Style -> Attr -> Attr
withoutStyle Style
Vty.bold
    Int
22 -> Style -> Attr -> Attr
withoutStyle Style
Vty.bold
    Int
23 -> Style -> Attr -> Attr
withoutStyle Style
Vty.standout
    Int
24 -> Style -> Attr -> Attr
withoutStyle Style
Vty.underline
    Int
25 -> Style -> Attr -> Attr
withoutStyle Style
Vty.blink
    Int
27 -> Style -> Attr -> Attr
withoutStyle Style
Vty.reverseVideo
    Int
i | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
30  Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
37   -> Color -> Attr -> Attr
withForeColor (Int -> Color
rawColor (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
30))
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
40  Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
47   -> Color -> Attr -> Attr
withBackColor (Int -> Color
rawColor (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
40))
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
90  Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
97   -> Color -> Attr -> Attr
withForeColor (Int -> Color
rawBrightColor (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
90))
      | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
100 Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
107  -> Color -> Attr -> Attr
withBackColor (Int -> Color
rawBrightColor (Int
i Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
100))
    Int
39 -> Attr -> Attr
resetForeColor
    Int
49 -> Attr -> Attr
resetBackColor
    Int
_  -> Attr -> Attr
forall a. a -> a
id
  where
    rawColor :: Int -> Color
rawColor = \case
        Int
0 -> Color
Vty.black
        Int
1 -> Color
Vty.red
        Int
2 -> Color
Vty.green
        Int
3 -> Color
Vty.yellow
        Int
4 -> Color
Vty.blue
        Int
5 -> Color
Vty.magenta
        Int
6 -> Color
Vty.cyan
        Int
_ -> Color
Vty.white
    rawBrightColor :: Int -> Color
rawBrightColor = \case
        Int
0 -> Color
Vty.brightBlack
        Int
1 -> Color
Vty.brightRed
        Int
2 -> Color
Vty.brightGreen
        Int
3 -> Color
Vty.brightYellow
        Int
4 -> Color
Vty.brightBlue
        Int
5 -> Color
Vty.brightMagenta
        Int
6 -> Color
Vty.brightCyan
        Int
_ -> Color
Vty.brightWhite
    withStyle :: Style -> Attr -> Attr
withStyle = (Attr -> Style -> Attr) -> Style -> Attr -> Attr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Attr -> Style -> Attr
Vty.withStyle
    withForeColor :: Color -> Attr -> Attr
withForeColor = (Attr -> Color -> Attr) -> Color -> Attr -> Attr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Attr -> Color -> Attr
Vty.withForeColor
    withBackColor :: Color -> Attr -> Attr
withBackColor = (Attr -> Color -> Attr) -> Color -> Attr -> Attr
forall a b c. (a -> b -> c) -> b -> a -> c
flip Attr -> Color -> Attr
Vty.withBackColor
    withoutStyle :: Style -> Attr -> Attr
withoutStyle Style
style Attr
attr = case Attr -> MaybeDefault Style
Vty.attrStyle Attr
attr of
        Vty.SetTo Style
oldStyle | Style
oldStyle Style -> Style -> Bool
`Vty.hasStyle` Style
style
                   -> Attr
attr { attrStyle :: MaybeDefault Style
Vty.attrStyle = Style -> MaybeDefault Style
forall v. (Eq v, Show v, Read v) => v -> MaybeDefault v
Vty.SetTo (Style
oldStyle Style -> Style -> Style
forall a. Bits a => a -> a -> a
.&. Style -> Style
forall a. Bits a => a -> a
complement Style
style) }
        MaybeDefault Style
_otherwise -> Attr
attr
    resetForeColor :: Attr -> Attr
resetForeColor Attr
attr = Attr
attr { attrForeColor :: MaybeDefault Color
Vty.attrForeColor = MaybeDefault Color
forall v. MaybeDefault v
Vty.KeepCurrent }
    resetBackColor :: Attr -> Attr
resetBackColor Attr
attr = Attr
attr { attrBackColor :: MaybeDefault Color
Vty.attrBackColor = MaybeDefault Color
forall v. MaybeDefault v
Vty.KeepCurrent }