{-
Copyright © 2017 Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>

This program is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation; either version 2 of the License, or
(at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program; if not, write to the Free Software
Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE OverloadedStrings #-}
{- |
   Module      : Text.Pandoc.Lua
   Copyright   : Copyright © 2017 Albert Krewinkel
   License     : GNU GPL, version 2 or above

   Maintainer  : Albert Krewinkel <tarleb+pandoc@moltkeplatz.de>
   Stability   : alpha

Pandoc lua utils.
-}
module Text.Pandoc.Lua ( runLuaFilter ) where

import Control.Monad ( (>=>), when )
import Control.Monad.Trans ( MonadIO(..) )
import Data.HashMap.Lazy ( HashMap )
import Data.Text ( Text, pack, unpack )
import Data.Text.Encoding ( decodeUtf8 )
import Scripting.Lua ( LuaState, StackValue(..) )
import Scripting.Lua.Aeson ( newstate )
import Text.Pandoc.Definition ( Block(..), Inline(..), Pandoc(..) )
import Text.Pandoc.Lua.PandocModule
import Text.Pandoc.Lua.StackInstances ()
import Text.Pandoc.Walk

import qualified Data.HashMap.Lazy as HashMap
import qualified Scripting.Lua as Lua

runLuaFilter :: (MonadIO m)
             => FilePath -> [String] -> Pandoc -> m Pandoc
runLuaFilter filterPath args pd = liftIO $ do
  lua <- newstate
  Lua.openlibs lua
  -- create table in registry to store filter functions
  Lua.push lua ("PANDOC_FILTER_FUNCTIONS"::String)
  Lua.newtable lua
  Lua.rawset lua Lua.registryindex
  -- store module in global "pandoc"
  pushPandocModule lua
  Lua.setglobal lua "pandoc"
  status <- Lua.loadfile lua filterPath
  if (status /= 0)
    then do
      luaErrMsg <- unpack . decodeUtf8 <$> Lua.tostring lua 1
      error luaErrMsg
    else do
      Lua.call lua 0 1
      Just luaFilters <- Lua.peek lua (-1)
      Lua.push lua (map pack args)
      Lua.setglobal lua "PandocParameters"
      doc <- runAll luaFilters >=> luaFilter lua "filter_doc" $ pd
      Lua.close lua
      return doc

runAll :: [LuaFilter] -> Pandoc -> IO Pandoc
runAll [] = return
runAll (x:xs) = walkMWithLuaFilter x >=> runAll xs

luaFilter :: Lua.LuaState -> String -> Pandoc -> IO Pandoc
luaFilter lua luaFn x = do
  fnExists <- isLuaFunction lua luaFn
  if fnExists
    then walkM (Lua.callfunc lua luaFn :: Pandoc -> IO Pandoc) x
    else return x

walkMWithLuaFilter :: LuaFilter -> Pandoc -> IO Pandoc
walkMWithLuaFilter (LuaFilter lua inlineFnMap blockFnMap docFnMap) =
  walkM (execInlineLuaFilter lua inlineFnMap) >=>
  walkM (execBlockLuaFilter  lua blockFnMap)  >=>
  walkM (execDocLuaFilter    lua docFnMap)

type InlineFunctionMap = HashMap Text (LuaFilterFunction Inline)
type BlockFunctionMap  = HashMap Text (LuaFilterFunction Block)
type DocFunctionMap    = HashMap Text (LuaFilterFunction Pandoc)
data LuaFilter =
  LuaFilter LuaState InlineFunctionMap BlockFunctionMap DocFunctionMap

newtype LuaFilterFunction a = LuaFilterFunction { functionIndex :: Int }

execDocLuaFilter :: LuaState
                 -> HashMap Text (LuaFilterFunction Pandoc)
                 -> Pandoc -> IO Pandoc
execDocLuaFilter lua fnMap x = do
  let docFnName = "Doc"
  case HashMap.lookup docFnName fnMap of
    Nothing -> return x
    Just fn -> runLuaFilterFunction lua fn x

execBlockLuaFilter :: LuaState
                   -> HashMap Text (LuaFilterFunction Block)
                   -> Block -> IO Block
execBlockLuaFilter lua fnMap x = do
  let filterOrId constr = case HashMap.lookup constr fnMap of
                            Nothing -> return x
                            Just fn -> runLuaFilterFunction lua fn x
  case x of
    Plain _          -> filterOrId "Plain"
    Para _           -> filterOrId "Para"
    LineBlock _      -> filterOrId "LineBlock"
    CodeBlock _ _    -> filterOrId "CodeBlock"
    RawBlock _ _     -> filterOrId "RawBlock"
    BlockQuote _     -> filterOrId "BlockQuote"
    OrderedList _ _  -> filterOrId "OrderedList"
    BulletList _     -> filterOrId "BulletList"
    DefinitionList _ -> filterOrId "DefinitionList"
    Header _ _ _     -> filterOrId "Header"
    HorizontalRule   -> filterOrId "HorizontalRule"
    Table _ _ _ _ _  -> filterOrId "Table"
    Div _ _          -> filterOrId "Div"
    Null             -> filterOrId "Null"

execInlineLuaFilter :: LuaState
                    -> HashMap Text (LuaFilterFunction Inline)
                    -> Inline -> IO Inline
execInlineLuaFilter lua fnMap x = do
  let runFn :: PushViaFilterFunction Inline a => LuaFilterFunction Inline -> a
      runFn fn = runLuaFilterFunction lua fn
  let tryFilter :: Text -> (LuaFilterFunction Inline -> IO Inline) -> IO Inline
      tryFilter fnName callFilterFn =
        case HashMap.lookup fnName fnMap of
          Nothing -> return x
          Just fn -> callFilterFn fn
  case x of
    LineBreak                 -> tryFilter "LineBreak" runFn
    SoftBreak                 -> tryFilter "SoftBreak" runFn
    Space                     -> tryFilter "Space"     runFn
    Cite cs lst               -> tryFilter "Cite"      $ \fn -> runFn fn lst cs
    Code attr str             -> tryFilter "Code"      $ \fn -> runFn fn str attr
    Emph lst                  -> tryFilter "Emph"      $ \fn -> runFn fn lst
    Math mt lst               -> tryFilter "Math"      $ \fn -> runFn fn lst mt
    Note blks                 -> tryFilter "Note"      $ \fn -> runFn fn blks
    Quoted qt lst             -> tryFilter "Quoted"    $ \fn -> runFn fn qt lst
    RawInline f str           -> tryFilter "RawInline" $ \fn -> runFn fn f str
    SmallCaps lst             -> tryFilter "SmallCaps" $ \fn -> runFn fn lst
    Span attr lst             -> tryFilter "Span"      $ \fn -> runFn fn lst attr
    Str str                   -> tryFilter "Str"       $ \fn -> runFn fn str
    Strikeout lst             -> tryFilter "Strikeout" $ \fn -> runFn fn lst
    Strong lst                -> tryFilter "Strong"    $ \fn -> runFn fn lst
    Subscript lst             -> tryFilter "Subscript" $ \fn -> runFn fn lst
    Superscript lst           -> tryFilter "Superscript" $ \fn -> runFn fn lst
    Link attr txt (src, tit)  -> tryFilter "Link" $
                                 \fn -> runFn fn txt src tit attr
    Image attr alt (src, tit) -> tryFilter "Image" $
                                 \fn -> runFn fn alt src tit attr

instance StackValue LuaFilter where
  valuetype _ = Lua.TTABLE
  push = undefined
  peek lua i = do
    -- TODO: find a more efficient way of doing this in a typesafe manner.
    inlineFnMap <- Lua.peek lua i
    blockFnMap  <- Lua.peek lua i
    docFnMap    <- Lua.peek lua i
    return $ LuaFilter lua <$> inlineFnMap <*> blockFnMap <*> docFnMap

-- | Helper class for pushing a single value to the stack via a lua function.
-- See @pushViaCall@.
class PushViaFilterFunction a b where
  pushViaFilterFunction' :: LuaState -> LuaFilterFunction a -> IO () -> Int -> b

instance (StackValue a) => PushViaFilterFunction a (IO a) where
  pushViaFilterFunction' lua lf pushArgs num = do
    pushFilterFunction lua lf
    pushArgs
    Lua.call lua num 1
    mbres <- Lua.peek lua (-1)
    case mbres of
      Nothing -> error $ "Error while trying to get a filter's return "
                 ++ "value from lua stack."
      Just res -> res <$ Lua.pop lua 1

instance (PushViaFilterFunction a c, StackValue b) =>
         PushViaFilterFunction a (b -> c) where
  pushViaFilterFunction' lua lf pushArgs num x =
    pushViaFilterFunction' lua lf (pushArgs *> push lua x) (num + 1)

-- | Push an value to the stack via a lua filter function. The function is
-- called with all arguments that are passed to this function and is expected to
-- return a single value.
runLuaFilterFunction :: (StackValue a, PushViaFilterFunction a b)
                     => LuaState -> LuaFilterFunction a -> b
runLuaFilterFunction lua lf = pushViaFilterFunction' lua lf (return ()) 0

-- | Push the filter function to the top of the stack.
pushFilterFunction :: Lua.LuaState -> LuaFilterFunction a -> IO ()
pushFilterFunction lua lf = do
  -- The function is stored in a lua registry table, retrieve it from there.
  push lua ("PANDOC_FILTER_FUNCTIONS"::String)
  Lua.rawget lua Lua.registryindex
  Lua.rawgeti lua (-1) (functionIndex lf)
  Lua.remove lua (-2) -- remove registry table from stack

instance StackValue (LuaFilterFunction a) where
  valuetype _ = Lua.TFUNCTION
  push lua v = pushFilterFunction lua v
  peek lua i = do
    isFn <- Lua.isfunction lua i
    when (not isFn) (error $ "Not a function at index " ++ (show i))
    Lua.pushvalue lua i
    push lua ("PANDOC_FILTER_FUNCTIONS"::String)
    Lua.rawget lua Lua.registryindex
    len <- Lua.objlen lua (-1)
    Lua.insert lua (-2)
    Lua.rawseti lua (-2) (len + 1)
    Lua.pop lua 1
    return . Just $ LuaFilterFunction (len + 1)


isLuaFunction :: Lua.LuaState -> String -> IO Bool
isLuaFunction lua fnName = do
  Lua.getglobal lua fnName
  res <- Lua.isfunction lua (-1)
  Lua.pop lua (-1)
  return res