{-# LANGUAGE OverloadedStrings #-}
module Text.Templating.Heist.Aeson
    ( JsonT
    , JsonTemplate
    , JsonState
    , JsonSpliceT
    , JsonSplice
    , addHeistAeson
    , runJsonT
    , renderJsonTemplate
    ) where

import Data.Aeson                            ( Value(..), Array )
import Text.Templating.Heist                 ( TemplateState, TemplateMonad
                                             , Splice, getParamNode, MIMEType
                                             , modifyTS, renderTemplate
                                             , runNodeList, bindSplices )
import Blaze.ByteString.Builder              ( Builder )
import qualified Data.Map as Map             ( lookup, toList )
import qualified Data.Vector as V            ( forM, toList )

import qualified Data.ByteString as Strict   ( ByteString )

import qualified Text.XmlHtml as X           ( Node(..), childNodes
                                             , getAttribute )
import qualified Data.Text as T              ( Text, append, pack, splitOn
                                             , intercalate )

import Control.Monad.Reader                  ( ReaderT, ask, local, runReaderT )


type JsonT m = ReaderT Value m
type JsonTemplate m a = TemplateMonad (JsonT m) a
type JsonState m = TemplateState (JsonT m)
type JsonSpliceT m = Splice (JsonT m)
type JsonSplice = JsonSpliceT IO

section :: Monad m => JsonSpliceT m
section = do
    node <- getParamNode
    withRequiredAttribute "key" $ \nodeIdentifier -> do
    withAnyValue nodeIdentifier $ \json -> do
    openJson json $ do
    runNodeList $ X.childNodes node

noSection :: Monad m => JsonSpliceT m
noSection = do
    node <- getParamNode
    withRequiredAttribute "key" $ \nodeIdentifier -> do
    json <- ask
    case findWithIdentifier nodeIdentifier json of
      Just{}  -> return []
      Nothing -> runNodeList $ X.childNodes node 

repeated :: Monad m => JsonSpliceT m
repeated = do
    node <- getParamNode
    withRequiredAttribute "key" $ \nodeIdentifier -> do
    withArrayValue nodeIdentifier $ \array -> do
    result <- V.forM array $ \entry -> openJson entry $ runNodeList $ X.childNodes node
    return $ concat (V.toList result)

withArrayValue ::Monad m => T.Text -> (Array -> JsonSpliceT m) -> JsonSpliceT m
withArrayValue nodeIdentifier action
    = withAnyValue nodeIdentifier $ \json ->
      case json of
        Array array -> action array
        _           -> return [X.TextNode $ "Json value expected to be an array: " `T.append` nodeIdentifier]

withAnyValue :: Monad m => T.Text -> (Value -> JsonSpliceT m) -> JsonSpliceT m
withAnyValue nodeIdentifier action
    = do json <- ask
         case findWithIdentifier nodeIdentifier json of
           Nothing    -> return []
           Just value -> action value

{-
  Return corresponding json values for identifiers such as:
    "@"
    "key"
    "object.key"
-}
findWithIdentifier :: T.Text -> Value -> Maybe Value
findWithIdentifier "@" currentJsonScope = Just currentJsonScope
findWithIdentifier identifier currentJsonScope
    = worker (T.splitOn "." identifier) currentJsonScope
    where worker [] json = Just json
          worker (key:keys) json
              = case json of
                  Object object | Just value <- Map.lookup key object
                    -> worker keys value
                  _ -> Nothing

withRequiredAttribute :: Monad m => T.Text -> (T.Text -> JsonSpliceT m) -> JsonSpliceT m
withRequiredAttribute attribute action
    = do node <- getParamNode
         case X.getAttribute attribute node of
           Nothing  -> return [X.TextNode ("Missing attribute: " `T.append` attribute)]
           Just val -> action val

-- Set a new json value as the current scope.
-- I'm not quite happy with the way it works. 2011-03-30
openJson :: Monad m => Value -> JsonTemplate m a -> JsonTemplate m a
openJson json splice
    = ask >>= \originalJsonScope ->
      local (const json) $
      do modifyTS $ closeJsonTS originalJsonScope
         modifyTS $ openJsonTS json
         tpl <- splice
         modifyTS $ closeJsonTS json
         return tpl

openJsonTS :: Monad m => Value -> TemplateState m -> TemplateState m
openJsonTS json
    = bindSplices [ ("val:" `T.append` key, return [X.TextNode val])
                    | (key, val) <- jsonToBinds json ]

-- Undo the binds created with 'openJsonTS'.
closeJsonTS :: Monad m => Value -> TemplateState m -> TemplateState m
closeJsonTS json
    = bindSplices [ ("val:" `T.append` key, generateOriginal)
                    | (key, _val) <- jsonToBinds json ]
    where generateOriginal = do node <- getParamNode
                                return [node]

jsonToBinds :: Value -> [(T.Text, T.Text)]
jsonToBinds = worker []
    where worker hierarchy json
              = case json of
                  Object object -> concat [ worker (key:hierarchy) val | (key, val) <- Map.toList object ]
                  Array{}       -> []
                  String string -> [(mkIndex hierarchy, string)]
                  Number num    -> [(mkIndex hierarchy, T.pack $ show num)]
                  Bool bool     -> [(mkIndex hierarchy, T.pack $ show bool)]
                  Null          -> [(mkIndex hierarchy, "Null")]
          mkIndex [] = "self"
          mkIndex hierarchy = T.intercalate "." (reverse hierarchy)

addHeistAeson :: Monad m => JsonState m -> JsonState m
addHeistAeson
    = bindSplices [ ("json:section", section)
                  , ("json:no-section", noSection)
                  , ("json:repeated", repeated) ]

renderJsonTemplate :: Monad m => JsonState m -> Strict.ByteString -> Value -> m (Maybe (Builder, MIMEType))
renderJsonTemplate state tplName json
    = runReaderT (renderTemplate (openJsonTS json state) tplName) json

runJsonT :: Monad m => JsonT m a -> Value -> m a
runJsonT action json
    = runReaderT action json