import Text.ParserCombinators.Parsec
import System.Environment
import System.IO
import System.Exit
import Control.Applicative ((<$>))
import Data.Char
import Data.Maybe
import Text.XHtml as X
import Control.Monad
import Control.Monad.State as S

-- the (Integer,Integer) state holds a (counter,columns) state
type MyParser = GenParser Char (Integer,Integer)

data ANSIEnv = AnsiEnv { ansiBrightOrBold :: Bool
                       , ansiFaint :: Bool
                       , ansiItalic :: Bool
                       , ansiUnderline :: Bool
                       , ansiReverse :: Bool
                       , ansiFontNumber :: Int
                       , ansiFramed :: Bool
                       , ansiEncircled :: Bool
                       , ansiOverlined :: Bool
                       , ansiFGColor :: ANSIColor
                       , ansiBGColor :: ANSIColor
                       }

defaultEnv :: ANSIEnv
defaultEnv = AnsiEnv False False False False False 0 False False False AnsiWhite AnsiBlack

modifyEnv :: ANSIEnv -> ANSIToken -> ANSIEnv
modifyEnv env ctl = if (finalByte ctl /= 'm')
    then env
    else foldl (flip ($)) env modifiers
 where modifiers = flip map (optParms ctl)
        (\parm -> case parm of
            0  -> \_ -> defaultEnv
            1  -> \e -> e { ansiBrightOrBold = True }
            2  -> \e -> e { ansiFaint = True }
            3  -> \e -> e { ansiItalic = True }
            4  -> \e -> e { ansiUnderline = True }
            7  -> \e -> e { ansiReverse = True }
            10 -> \e -> e { ansiFontNumber = 0 }
            11 -> \e -> e { ansiFontNumber = 1 }
            12 -> \e -> e { ansiFontNumber = 2 }
            13 -> \e -> e { ansiFontNumber = 3 }
            14 -> \e -> e { ansiFontNumber = 4 }
            15 -> \e -> e { ansiFontNumber = 5 }
            16 -> \e -> e { ansiFontNumber = 6 }
            17 -> \e -> e { ansiFontNumber = 7 }
            18 -> \e -> e { ansiFontNumber = 8 }
            19 -> \e -> e { ansiFontNumber = 9 }
            21 -> \e -> e { ansiBrightOrBold = False }
            22 -> \e -> e { ansiBrightOrBold = False
                          , ansiFaint = False
                          }
            23 -> \e -> e { ansiItalic = False }
            24 -> \e -> e { ansiUnderline = False }
            26 -> \e -> e { ansiUnderline = False }
            27 -> \e -> e { ansiReverse = False }
            30 -> \e -> e { ansiFGColor = AnsiBlack }
            31 -> \e -> e { ansiFGColor = AnsiRed }
            32 -> \e -> e { ansiFGColor = AnsiGreen }
            33 -> \e -> e { ansiFGColor = AnsiYellow }
            34 -> \e -> e { ansiFGColor = AnsiBlue }
            35 -> \e -> e { ansiFGColor = AnsiMagenta }
            36 -> \e -> e { ansiFGColor = AnsiCyan }
            37 -> \e -> e { ansiFGColor = AnsiWhite }
            39 -> \e -> e { ansiFGColor = AnsiWhite }
            40 -> \e -> e { ansiBGColor = AnsiBlack }
            41 -> \e -> e { ansiBGColor = AnsiRed }
            42 -> \e -> e { ansiBGColor = AnsiGreen }
            43 -> \e -> e { ansiBGColor = AnsiYellow }
            44 -> \e -> e { ansiBGColor = AnsiBlue }
            45 -> \e -> e { ansiBGColor = AnsiMagenta }
            46 -> \e -> e { ansiBGColor = AnsiCyan }
            47 -> \e -> e { ansiBGColor = AnsiWhite }
            49 -> \e -> e { ansiBGColor = AnsiBlack }
            51 -> \e -> e { ansiFramed = True }
            52 -> \e -> e { ansiEncircled = True }
            53 -> \e -> e { ansiOverlined = True }
            54 -> \e -> e { ansiFramed = False
                          , ansiEncircled = False
                          }
            55 -> \e -> e { ansiOverlined = False }
            _  -> id)


data ANSIColor = AnsiBlack
               | AnsiRed
               | AnsiGreen
               | AnsiYellow
               | AnsiBlue
               | AnsiMagenta
               | AnsiCyan
               | AnsiWhite
    deriving Eq

envAttrs :: ANSIEnv -> [HtmlAttr]
envAttrs env =
    let
        bob       = when' (ansiBrightOrBold env) ("bold")
        italic    = when' (ansiItalic env) ("italic")
        underline = when' (ansiUnderline env) ("underline")
        fontnum   = when' (ansiFontNumber env /= 0) ("fontNum" ++ (show $ ansiFontNumber env))
        framed    = when' (ansiFramed env) ("framed")
        fgcolor   = when' (True) (reversed "bg" "fg" ++ (show $ ansiFGColor env))
        bgcolor   = when' (True) (reversed "fg" "bg" ++ (show $ ansiBGColor env))
        classes = map ("ansi2html-"++) $ catMaybes [bob, italic, underline, fontnum, framed, fgcolor, bgcolor]
    in [theclass $ unwords classes]
    where
        reversed a b = if ansiReverse env then a else b
        when' x y = if x then Just y else Nothing

instance Show ANSIColor where
    show color = case color of
        AnsiBlack       -> "black"
        AnsiRed         -> "red"
        AnsiGreen       -> "green"
        AnsiYellow      -> "yellow"
        AnsiBlue        -> "blue"
        AnsiMagenta     -> "magenta"
        AnsiCyan        -> "cyan"
        AnsiWhite       -> "white"

data ANSIToken = AnsiText String
           | CtlSeq { privModeChars :: String
                    , optParms      :: [Int]
                    , interMedChars :: String
                    , finalByte     :: Char
                    }
    deriving Show

ansiParser :: MyParser [ANSIToken]
ansiParser = do
    result <- many ansiToken
    eof
    return result

ansiToken :: MyParser ANSIToken
ansiToken = do
    {- xterm prints a strange file with each line prefixed by '\ESC#5'
     - Let's simply remove those prefixes -}
    try (string "\ESC#5") <|> return ""
    try ctlSeq <|> ansiText

ansiText :: MyParser ANSIToken
ansiText = do
    strings <- many1 nonControlChar
    column <- getState
    return $ AnsiText $ concat strings

nonControlChar :: MyParser String
nonControlChar = do
    -- filter out \CR
    try (char '\CR') <|> return ' '
    (counter,columns) <- getState
    r <- satisfy (/='\ESC')
    implicitbreak <- if counter == columns && r /= '\n'
                        then do
                            updateState (mapFst (const 0))
                            return "\n"
                        else return ""
    if (r == '\n')
     then do
        updateState (mapFst (const 0))
     else do
        updateState (mapFst (+1))
    return $ implicitbreak ++ [r]
    where mapFst f (a,b) = (f a,b)

ctlSeq :: MyParser ANSIToken
ctlSeq = do
    csiCode
    parms <- number `sepBy` char ';'
    iMBs <- interMedBytes
    l <- letter
    return $ CtlSeq "" parms iMBs l

number  :: MyParser Int
number  = do{ ds <- many1 digit
            ; return (read ds)
            }
        <?> "number"

interMedBytes :: MyParser String
interMedBytes = return ""

csiCode :: MyParser ()
csiCode = do
    escapeChar
    char '['
    return ()

escapeChar :: MyParser ()
escapeChar = do
    char '\ESC'
    return ()

ansiToken2Html :: ANSIToken -> S.State ANSIEnv X.Html
ansiToken2Html tok = do
    env <- S.get
    case tok of
        AnsiText t -> let html = (linesToHtml $ lines t) +++ (if last t == '\n' then X.br else X.noHtml)
                        in return $ (X.thespan ! envAttrs env) << html
        ctlseq     -> do
                        S.put $ modifyEnv env ctlseq
                        return X.noHtml

ansi2Html :: [ANSIToken] -> X.Html
ansi2Html toks = foldl (+++) noHtml htmls
    where
        htmls = evalState (mapM ansiToken2Html toks) defaultEnv

makeHtml :: X.Html -> X.Html
makeHtml content = X.thediv ! [theclass "ansi2html"] << content

printUsage :: IO ()
printUsage = do
    hPutStrLn stderr "Usage: ansi2html <columns>"
    hFlush stderr

main = do
    args <- getArgs
    columns <- case args of
        []    -> return 80
        x:xs  -> case reads x :: [(Integer,String)] of
                    [(c,[])]    -> return c
                    _           -> printUsage >> exitFailure
    contents <- getContents
    tokens <- case runParser ansiParser (0,columns) "" contents of
                Left err -> do
                        error $ "parse error at " ++ show err
                Right x -> return x
    print . makeHtml $ ansi2Html tokens +++ dummyLine columns
    where
    -- force width of terminal. Is there a better solution?
    dummyLine columns = X.thespan ! [X.theclass "ansi2html-dummyline"] <<
                    (X.lineToHtml $ replicate (fromInteger columns) ' ')