{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UnicodeSyntax #-}
module Main where

import Control.Applicative ((<$>), (<|>), many)
import Data.Monoid (Monoid(..), (<>))

import           Control.Lens (makeLenses, over, view)
import           Data.Text.Lazy (Text)
import qualified Data.Text.Lazy    as T
import qualified Data.Text.Lazy.IO as T
import           Options.Applicative (execParser, fullDesc, header, helper, info, progDesc)
import           Text.Parsec (anyChar, parse, string, try)
import           Text.Parsec.Text.Lazy


data Chunk =
    Raw { _raw ∷ Text }
  | Ok { _raw ∷ Text }
  | Fail { _raw ∷ Text }
  | Error { _raw ∷ Text }
    deriving (Show, Read)


makeLenses ''Chunk


main ∷ IO ()
main = do
  execParser opts
  T.interact (T.unlines . map line . T.lines)
 where
  opts = info helper
    ( fullDesc
    <> progDesc "Colorize python unittest execution output"
    <> header "Hpyrg - Pyrg utility done right"
    )
  line = wrap . mconcat . map colorize . T.words


colorize ∷ Text → Chunk
colorize input = case parse word "(hpyrg)" input of
  Right w → w
  Left _ → Raw input


instance Monoid Chunk where
  mempty = Raw mempty
  Raw r `mappend` w      = over raw (\s → r <> " " <> s) w
  Ok l `mappend` Raw r   = over raw (\s → s <> " " <> r) (Ok l)
  Ok r `mappend` w       = over raw (\s → r <> " " <> s) w
  Fail l `mappend` Raw r = over raw (\s → s <> " " <> r) (Fail l)
  Fail l `mappend` Ok r  = over raw (\s → s <> " " <> r) (Fail l)
  Fail r `mappend` w     = over raw (\s → r <> " " <> s) w
  Error r `mappend` w    = over raw (\s → s <> " " <> view raw w) (Error r)


wrap ∷ Chunk → Text
wrap (Raw t)   = t
wrap (Ok t)    = "\ESC[32m"   <> t <> "\ESC[0m"
wrap (Fail t)  = "\ESC[1;33m" <> t <> "\ESC[0m"
wrap (Error t) = "\ESC[31m"   <> t <> "\ESC[0m"


word ∷ Parser Chunk
word =
  Ok . T.pack    <$> try (string "ok")    <|>
  Ok . T.pack    <$> try (string "OK")    <|>
  Fail . T.pack  <$> try (string "FAIL")  <|>
  Error . T.pack <$> try (string "ERROR") <|>
  Raw . T.pack   <$> many anyChar