{-# LANGUAGE Arrows #-}

module Bamboo.Plugin.Highlight (highlight) where


import Text.XML.HXT.Arrow hiding (first, mkText)
import Text.XML.HXT.Parser.HtmlParsec
import MPSUTF8
import Prelude hiding ((.), (>))
import Text.Highlighting.Kate
import Text.XHtml.Strict hiding (body)
import Hack
import Hack.Contrib.Middleware.Censor
import Data.ByteString.Lazy.Char8 (pack, unpack)

highlight :: Middleware
highlight = censor code
  where
    code r = do
      bodies <- r.body.unpack.play
      let b = if bodies.null then r.body else bodies.first.pack
      return r {body = b}

play :: String -> IO [String]
play s = runX (processor s)

processor :: String -> IOSArrow XmlTree String
processor s =
  readString 
    [ (a_validate,"0")
    , (a_parse_html, "1")
    , (a_encoding, utf8)
    ] s
  > process_source
  > writeDocumentToString 
    [ (a_output_encoding, utf8)
    , (a_output_html, "1")
    , (a_no_empty_elements, "1")
    ]
  where
    process_source = processBottomUp (highlight_code `when` is_code_block)
       where
        highlight_code = replaceChildren highlighted 
          > changeElemName (const $ mkName "div")
      
        highlighted = proc x -> do
          code <- getChildren > getChildren > getText -< x
          language <- getAttrValue "class" -< x
          let h = kate language code
          returnA -< parseHtmlContent h .first
      
        is_code_block = 
            isElem
          > hasName "pre"
          > hasAttr "class"
          > getChildren
          > isElem
          > hasName "code"
  
    kate language code = 
      case highlightAs language code of
        Right result -> renderHtmlFragment $ 
          formatAsXHtml [OptNumberLines] language result
        Left  _    -> "<pre><code>" ++ code ++ "</pre></code>"