{-# OPTIONS -fno-warn-orphans #-}
module Darcs.ColorPrinter ( errorDoc, traceDoc, assertDoc, fancyPrinters ) where

import Debug.Trace ( trace )
import System.IO ( stderr )
import Darcs.External (getTermNColors)
import Printer (Printer, Printers, Printers'(..), Printable(..), Color(..),
                invisiblePrinter, (<>), (<?>), Doc(Doc,unDoc), unsafeBothText, simplePrinter, hcat,
                unsafeText, unsafeChar, space, unsafePackedString,
                renderStringWith, prefix )
import Data.Char ( isAscii, isPrint, isSpace, isControl, ord, chr, intToDigit )
import Data.Bits ( bit, xor )
import System.Environment ( getEnv )
import qualified Data.ByteString.Char8 as BC (unpack, any, last, spanEnd)
import qualified Data.ByteString       as B (null, init)
import System.IO.Unsafe ( unsafePerformIO )
import System.IO ( hIsTerminalDevice, Handle )

dollar, cr :: Doc
dollar = unsafeBothText "$"
cr     = unsafeBothText "\r"

errorDoc :: Doc -> a
errorDoc = error . show

traceDoc :: Doc -> a -> a
traceDoc d = trace (show d)

assertDoc :: Maybe Doc -> a -> a
assertDoc Nothing x = x
assertDoc (Just e) _ = errorDoc e

instance Show Doc where
    show = renderStringWith (fancyPrinters stderr)

-- policy
-- | the 'Policy' type is a record containing the variables which control
-- how 'Doc's will be rendered on some output.
data Policy = Policy { poColor :: Bool    -- ^ overall use of color
                     , poEscape :: Bool   -- ^ overall use of escaping
                     , poLineColor :: Bool -- ^ overall use of colored lines (only hunks for now)
                     , poAltColor :: Bool -- ^ alternative to color (bold, inverse)
                     , poIsprint :: Bool  -- ^ don't escape isprints
                     , po8bit  :: Bool    -- ^ don't escape 8-bit chars
                     , poNoEscX :: String   -- ^ extra chars to never escape
                     , poEscX :: String   -- ^ extra chars to always escape
                     , poTrailing :: Bool -- ^ escape trailing spaces
                     , poCR :: Bool       -- ^ ignore \r at end of lines
                     , poSpace :: Bool    -- ^ escape spaces (used with poTrailing)
                     }

{-# NOINLINE getPolicy #-}
-- | 'getPolicy' returns a suitable policy for a given handle.
-- The policy is chosen according to environment variables, and to the
-- type of terminal which the handle represents
getPolicy :: Handle -> Policy
getPolicy handle = unsafePerformIO $
 do isTerminal <- hIsTerminalDevice handle
    nColors <- if isTerminal then getTermNColors else return 0

    envDontEscapeAnything  <- getEnvBool "DARCS_DONT_ESCAPE_ANYTHING"
    envDontEscapeIsprint   <- getEnvBool "DARCS_DONT_ESCAPE_ISPRINT"
    envUseIsprint          <- getEnvBool "DARCS_USE_ISPRINT" -- for backwards-compatibility
    envDontEscape8bit      <- getEnvBool "DARCS_DONT_ESCAPE_8BIT"

    envDontEscapeExtra  <- getEnvString "DARCS_DONT_ESCAPE_EXTRA"
    envEscapeExtra      <- getEnvString "DARCS_ESCAPE_EXTRA"

    envDontEscapeTrailingSpace  <- getEnvBool "DARCS_DONT_ESCAPE_TRAILING_SPACES"
    envDontEscapeTrailingCR     <- getEnvBool "DARCS_DONT_ESCAPE_TRAILING_CR"

    envDontColor         <- getEnvBool "DARCS_DONT_COLOR"
    envAlwaysColor       <- getEnvBool "DARCS_ALWAYS_COLOR"
    envAlternativeColor  <- getEnvBool "DARCS_ALTERNATIVE_COLOR"
    envDoColorLines    <- getEnvBool "DARCS_DO_COLOR_LINES"

    let haveColor = envAlwaysColor || (isTerminal && (nColors > 4))
        doColor   = not envDontColor && haveColor

    return Policy { poColor    = doColor,
                    poEscape   = not envDontEscapeAnything,
                    poLineColor= doColor && envDoColorLines,
                    poIsprint  = envDontEscapeIsprint || envUseIsprint,
                    po8bit     = envDontEscape8bit,
                    poNoEscX   = envDontEscapeExtra,
                    poEscX     = envEscapeExtra,
                    poTrailing = not envDontEscapeTrailingSpace,
                    poCR       = envDontEscapeTrailingCR,
                    poAltColor = haveColor && envAlternativeColor,

                    poSpace = False
                  }
 where
  getEnvBool s = safeGetEnv s >>= return.(/= "0")
  safeGetEnv s = getEnv s `catch` \_ -> return "0"
  getEnvString s = getEnv s `catch` \_ -> return ""


-- printers

-- | @'fancyPrinters' h@ returns a set of printers suitable for outputting
-- to @h@
fancyPrinters :: Printers
fancyPrinters h = let policy = getPolicy h in 
                      Printers { colorP = colorPrinter policy,
                             invisibleP = invisiblePrinter,
                             hiddenP = colorPrinter policy Green,
                             userchunkP = userchunkPrinter policy,
                             defP       = escapePrinter policy,
                             lineColorT = lineColorTrans policy,
                             lineColorS = lineColorSuffix policy
                           }

-- | @'lineColorTrans' policy@ tries to color a Doc, according to policy po.
-- That is, if @policy@ has @poLineColor@ set, then colors the line, otherwise
-- does nothing.
lineColorTrans :: Policy -> Color -> Doc -> Doc
lineColorTrans po | poLineColor po = \c d -> prefix (set_color c) d <?> unsafeBothText reset_color
                  | otherwise      = const id

lineColorSuffix :: Policy -> [Printable] -> [Printable]
lineColorSuffix po | poLineColor po = \d -> S reset_color : d
                   | otherwise      = id

colorPrinter :: Policy -> Color -> Printer
colorPrinter po | poColor po = \c -> unDoc . color po c . Doc . escapePrinter po{poColor=False}
                | otherwise  = const $ escapePrinter po

userchunkPrinter :: Policy -> Printer
userchunkPrinter po p
 | not (poEscape po)   = simplePrinter p
 | not (poTrailing po) = escapePrinter po p
 | otherwise           = unDoc $ pr p
 where
  pr (S s)       = prString s
  pr (Both _ ps) = prPS ps
  pr (PS ps)     = prPS ps

  prPS ps = let (leadPS, trailPS) = BC.spanEnd isSpace ps
            in if B.null trailPS
                then Doc $ escapePrinter po p
                else Doc (escapePrinter po (PS leadPS))
                  <> Doc (escapePrinter po{poSpace=True} (PS trailPS))
                  <> mark_escape po dollar

  prString s = let (trail',lead') = span isSpace (reverse s)
                   lead = reverse lead'
                   trail = reverse trail'
               in if (not.null) trail
                   then Doc (escapePrinter po (S lead))
                     <> Doc (escapePrinter po{poSpace=True} (S trail))
                     <> mark_escape po dollar
                   else Doc (escapePrinter po p)

escapePrinter :: Policy -> Printer
escapePrinter po
 | (not.poEscape) po = simplePrinter
 | otherwise         = unDoc . crepr
 where
  crepr p | poCR po && isEndCR p = epr (initPR p) <> cr
          | otherwise            = epr p

  epr (S s)      = escape po s
  epr (PS ps)    = if BC.any (not.no_escape po) ps
                   then escape po (BC.unpack ps)
                   else unsafePackedString ps
  epr (Both s _) = escape po s

  isEndCR (S s)        = not (null s) && last s == '\r'
  isEndCR (PS ps)      = not (B.null ps) && BC.last ps == '\r'
  isEndCR (Both _ ps)  = not (B.null ps) && BC.last ps == '\r'

  initPR (S s)       = S $ init s
  initPR (PS ps)     = PS $ B.init ps
  initPR (Both s ps) = Both (init s) (B.init ps)


-- escape assumes the input is in ['\0'..'\255']

-- | @'escape' policy string@ escapes @string@ according to the rules
-- defined in 'policy', turning it into a 'Doc'.
escape :: Policy -> String -> Doc
escape _ "" = unsafeText ""
escape po s = hcat (map escapeChar s)
 where
  escapeChar c | no_escape po c = unsafeChar c
  escapeChar ' ' = space
  escapeChar c = (emph.unsafeText.quoteChar) c
  emph = mark_escape po

-- | @'no_escape' policy c@ tells wether @c@ will be left as-is
-- when escaping according to @policy@
no_escape :: Policy -> Char -> Bool
no_escape po c | poSpace po && isSpace c = False
no_escape po c | c `elem` poEscX po = False
no_escape po c | c `elem` poNoEscX po = True
no_escape _ '\t' = True  -- tabs will likely be converted to spaces
no_escape _ '\n' = True
no_escape po c = if (poIsprint po) then isPrint c
                                   else isPrintableAscii c
                 ||  c >= '\x80' && po8bit po

-- | 'isPrintableAscii' tells wether a character is a printable character
-- of the ascii range.
isPrintableAscii :: Char -> Bool
isPrintableAscii c = isAscii c && isPrint c


-- | 'quoteChar' represents a special character as a string.
--   * @quoteChar '^c'@ (where @^c@ is a control character) is @"^c"@
--   * Otherwise, @quoteChar@ returns "\hex", where 'hex' is the
--     hexadecimal number of the character.
quoteChar :: Char -> String
quoteChar c
 | isControl c && isPrintableAscii cHat = ['^', cHat]
 | otherwise = sHex
 where
  cHat = chr $ (bit 6 `xor`) $ ord c
  sHex = let (q, r) = quotRem (ord c) 16
         in ['\\', intToDigit q, intToDigit r]


-- make colors and highlightings

-- | @'mark_escape' policy doc@ marks @doc@ with the appropriate
-- marking for escaped characters according to @policy@
mark_escape :: Policy -> Doc -> Doc
mark_escape po | poAltColor po  = make_invert
               | poColor po     = make_color Red
               | otherwise      = make_asciiart

-- | @'color' policy color doc@ colors @doc@ with color @color@ if
-- @policy@ is not set to use an alternative to color. In that case,
-- it makes the text bold instead.
color :: Policy -> Color -> Doc -> Doc
color po | poAltColor po = \_ -> make_bold
         | otherwise     = make_color

make_color, make_color' :: Color -> Doc -> Doc

make_color' = with_color . set_color

-- memoized version of make_color'
make_color Blue    = make_color' Blue
make_color Red     = make_color' Red
make_color Green   = make_color' Green
make_color Cyan    = make_color' Cyan
make_color Magenta = make_color' Magenta

set_color :: Color -> String
set_color Blue    = "\x1B[01;34m" -- bold blue
set_color Red     = "\x1B[01;31m" -- bold red
set_color Green   = "\x1B[01;32m" -- bold green
set_color Cyan    = "\x1B[36m"    -- light cyan
set_color Magenta = "\x1B[35m"    -- light magenta

-- | @'make_asciiart' doc@ tries to make @doc@ (usually a
-- single escaped char) stand out with the help of only plain
-- ascii, i.e., no color or font style.
make_asciiart :: Doc -> Doc
make_asciiart x = unsafeBothText "[_" <> x <> unsafeBothText "_]"

-- | the string to reset the terminal's color.
reset_color :: String
reset_color = "\x1B[00m"

-- | @'with_color' color doc@ returns a colorized version of @doc@.
-- @color@ is a string that represents a color, given by 'set_color'
with_color :: String -> Doc -> Doc
with_color c =
   let c' = unsafeBothText c
       r' = unsafeBothText reset_color
   in \x -> c' <> x <> r'


-- | 'make_bold' boldens a doc.
make_bold :: Doc -> Doc
-- | 'make_invert' returns an invert video version of a doc.
make_invert :: Doc -> Doc
make_bold   = with_color "\x1B[01m"
make_invert = with_color "\x1B[07m"