\begin{code}
module ColourPrinter ( fancyPrinters ) where
import External (getTermNColors)
import Printer (Printer, Printers, Printers'(..), Printable(..), Color(..),
                invisiblePrinter, (<>), Doc, unsafeBoth, simplePrinter, hcat,
                unsafeText, unsafeChar, space, unsafePackedString)
import Char ( isAscii, isPrint, isSpace, isControl, ord, chr, intToDigit )
import Data.Bits ( bit, xor )
import System ( getEnv )
import FastPackedString ( packString, unpackPS, anyPS, lastPS, initPS,
                          nullPS, spanEndPS )
import System.IO.Unsafe ( unsafePerformIO )
import System.IO ( hIsTerminalDevice, Handle )

dollar, cr :: Doc
dollar = unsafeBoth  "$" (packString  "$")
cr     = unsafeBoth "\r" (packString "\r")


-- policy

data Policy = Policy { poColor,    -- overall use of color
                       poEscape,   -- overall use of escaping
                       poAltColor, -- alternative to color (bold, inverse)
                       poIsprint,  -- don't escape isprints
                       po8bit      -- don't escape 8-bit chars
                        :: Bool,
                       poNoEscX,   -- extra chars to never escape
                       poEscX      -- extra chars to always escape
                        :: String,
                       poTrailing, -- escape trailing spaces
                       poCR,       -- ignore \r at end of lines

                       poSpace         -- escape spaces (used with poTrailing)
                        :: Bool
                     }

{-# NOINLINE getPolicy #-}
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"

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

    return Policy { poColor    = not envDontColor && haveColor,
                    poEscape   = not envDontEscapeAnything,
                    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 :: Printers
fancyPrinters h = Printers { colorP     = colorPrinter (getPolicy h),
                             invisibleP = invisiblePrinter,
                             userchunkP  = userchunkPrinter (getPolicy h),
                             defP       = escapePrinter (getPolicy h)
                           }

colorPrinter :: Policy -> Color -> Printer
colorPrinter po = if poColor po
                   then  \c -> color po c . escapePrinter po{poColor=False}
                   else  \_ -> escapePrinter po

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

  prPS ps = let (leadPS, trailPS) = spanEndPS isSpace ps
            in if nullPS trailPS
                then escapePrinter po p
                else escapePrinter po (PS leadPS)
                  <> 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 escapePrinter po (S lead)
                     <> escapePrinter po{poSpace=True} (S trail)
                     <> mark_escape po dollar
                   else escapePrinter po p

escapePrinter :: Policy -> Printer
escapePrinter po
 | (not.poEscape) po = simplePrinter
 | otherwise         = 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 anyPS (not.no_escape po) ps
                   then escape po (unpackPS ps)
                   else unsafePackedString ps
  epr (Both s _) = escape po s

  isEndCR (S s)        = not (null s) && last s == '\r'
  isEndCR (PS ps)      = not (nullPS ps) && lastPS ps == '\r'
  isEndCR (Both _ ps)  = not (nullPS ps) && lastPS ps == '\r'

  initPR (S s)       = S $ init s
  initPR (PS ps)     = PS $ initPS ps
  initPR (Both s ps) = Both (init s) (initPS ps)


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

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 -> 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 :: Char -> Bool
isPrintableAscii c = isAscii c && isPrint c


-- quote

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 colours and highlightings

mark_escape :: Policy -> Doc -> Doc
mark_escape po | poAltColor po  = make_invert
               | poColor po     = make_red
               | otherwise      = make_asciiart

color :: Policy -> Color -> Doc -> Doc
color po | poAltColor po = \_ -> make_bold
         | otherwise     = make_color

make_color :: Color -> Doc -> Doc
make_color Blue  = make_blue
make_color Red   = make_red
make_color Green = make_green

make_asciiart :: Doc -> Doc
make_asciiart x = unsafeBoth "[_" (packString "[_")
               <> x
               <> unsafeBoth "_]" (packString "_]")

make_bold :: Doc -> Doc
make_bold x = unsafeBoth "\x1B[01m" (packString "\x1B[01m")
              <> x
              <> reset_colour

make_invert :: Doc -> Doc
make_invert x = unsafeBoth "\x1B[07m" (packString "\x1B[07m")
                <> x
                <> reset_colour

make_blue :: Doc -> Doc
make_blue x = unsafeBoth "\x1B[01;34m" (packString "\x1B[01;34m")
           <> x
           <> reset_colour

make_red :: Doc -> Doc
make_red x = unsafeBoth "\x1B[01;31m" (packString "\x1B[01;31m")
          <> x
          <> reset_colour

make_green :: Doc -> Doc
make_green x = unsafeBoth "\x1B[01;32m" (packString "\x1B[01;32m")
            <> x
            <> reset_colour

reset_colour :: Doc
reset_colour = unsafeBoth "\x1B[00m" (packString "\x1B[00m")

\end{code}