module Language.Haskell.HBB.Internal.TTreeColor (
    applyColoredTTree
    ) where

import Language.Haskell.HBB.Internal.SrcSpan
import Language.Haskell.HBB.Internal.TTree
import Control.Monad.State.Lazy
import Data.List (sortBy)
import SrcLoc

data TTreeColor = Gray
                | Red      | HlRed     | Green    | HlGreen
                | Yellow   | HlYellow  | Blue     | HlBlue
                | Magenta  | HlMagenta | Cyan     | HlCyan
                | DefaultWhiteOnBlack 

additionColor = Gray
displayColors = cycle [   Red ,   Green ,   Yellow ,   Blue ,   Magenta ,   Cyan
                      , HlRed , HlGreen , HlYellow , HlBlue , HlMagenta , HlCyan ]

ansiColorStr :: TTreeColor -> String
ansiColorStr Gray      = "\ESC[1;30m"
ansiColorStr Red       = "\ESC[0;31m"
ansiColorStr Green     = "\ESC[0;32m"
ansiColorStr Yellow    = "\ESC[0;33m"
ansiColorStr Blue      = "\ESC[0;34m"
ansiColorStr Magenta   = "\ESC[0;35m"
ansiColorStr Cyan      = "\ESC[0;36m"

ansiColorStr HlRed     = "\ESC[1;31m"
ansiColorStr HlGreen   = "\ESC[1;32m"
ansiColorStr HlYellow  = "\ESC[1;33m"
ansiColorStr HlBlue    = "\ESC[1;34m"
ansiColorStr HlMagenta = "\ESC[1;35m"
ansiColorStr HlCyan    = "\ESC[1;36m"

ansiColorStr DefaultWhiteOnBlack   = "\ESC[0m"

attachColors 
    :: (BufSpan,ClientTTree) 
    -> ((BufSpan,TTreeColor),TTree LineBuf (RealSrcSpan,Int) (BufSpan,TTreeColor))
attachColors tree = 

    -- This is the axiliary function that transforms a tree into a colored one.
    -- All additions will be gray and displays will use the rest of the colors
    -- (except white which is used for the rest of the text).
    let colorStep 
            :: (BufSpan,TTree LineBuf (RealSrcSpan,Int) BufSpan) 
            -> State (Int,[(RealSrcSpan,TTreeColor)]) ((BufSpan,TTreeColor),TTree LineBuf (RealSrcSpan,Int) (BufSpan,TTreeColor))
        colorStep (l,(TTree content childs)) = do
            childs' <- mapM colorStep childs
            color <- case content of
                (Display (spn,_)) -> do
                    -- Ok, this is a display.
                    (colorIdx,alreadyKnown) <- get

                    let (selectedColor,newState) = case filter (\(s,_) -> s == spn) alreadyKnown of
                            []      -> let color2add = (displayColors !! colorIdx)
                                       in  (color2add,(colorIdx+1,((spn,color2add):alreadyKnown)))
                            [(_,c)] -> (c,(colorIdx,alreadyKnown))
                            (_:_)   -> error "internal error (more than one color for one src-span)"
                    put newState
                    return selectedColor
                (Addition _) -> return additionColor
            return $ ((l,color),(TTree content childs'))

    in evalState (colorStep tree) (0,[])

-- This function is the pendant to 'applyTTree' from
-- 'Language.Haskell.HBB.Internal.TTree'.  It produces a TTree which contains ANSI
-- Escape sequences for a colored output in an according terminal.
--
-- Text that has not been altered by the transformation-tree is written with
-- the default settings. Additions are written in gray and Displays use the
-- rest of the color-space.
--
-- This function makes use of the function applyTTreeGeneric which has
-- intentionally been made as generic as needed to be used here.
applyColoredTTree :: 
    [(FilePath,LineBuf)]     -- File cache 
    -> (BufSpan,ClientTTree) -- The tree itself
    -> LineBuf               -- The file to transform
    -> LineBuf
applyColoredTTree fc tree lns = 
    applyTTreeGeneric
        (Just (DefaultWhiteOnBlack,(\(_,x) -> x),finFun)) 
        fc
        (\(s,_) -> s)
        (attachColors tree)
        lns

    where
        finFun :: (TTreeColor,(BufSpan,TTreeColor)) -> (LineBuf,LineBuf,LineBuf) -> (LineBuf,LineBuf,LineBuf)
        finFun (parColor,(_,color)) (initLines,childsRes,traiLines) = 
            (                                   initLines
            ,joinSplit ([ansiColorStr    color],childsRes)
            ,joinSplit ([ansiColorStr parColor],traiLines))