1 files changed, 111 insertions, 14 deletions

diff --git a/src/Text/Pandoc/Lua/PandocModule.hs b/src/Text/Pandoc/Lua/PandocModule.hs
index afb9aeca6..ffd681d30 100644
--- a/src/Text/Pandoc/Lua/PandocModule.hs
+++ b/src/Text/Pandoc/Lua/PandocModule.hs
@@ -15,6 +15,10 @@ 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 CPP #-}
+#if !MIN_VERSION_hslua(0,9,0)
+{-# OPTIONS_GHC -fno-warn-orphans #-}
+#endif
 {- |
    Module      : Text.Pandoc.Lua.PandocModule
    Copyright   : Copyright © 2017 Albert Krewinkel
@@ -25,28 +29,37 @@ Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
 
 Pandoc module for lua.
 -}
-module Text.Pandoc.Lua.PandocModule ( pushPandocModule ) where
+module Text.Pandoc.Lua.PandocModule
+  ( pushPandocModule
+  , pushMediaBagModule
+  ) where
 
-import Control.Monad (unless)
+import Control.Monad (unless, zipWithM_)
 import Data.ByteString.Char8 (unpack)
 import Data.Default (Default (..))
+import Data.IORef
 import Data.Text (pack)
-import Foreign.Lua (Lua, Status (OK), NumResults, call, loadstring, liftIO,
-                    push, pushHaskellFunction, rawset)
-import Text.Pandoc.Class (readDataFile, runIO, runIOorExplode, setUserDataDir)
+import Foreign.Lua (Lua, FromLuaStack, ToLuaStack, NumResults, liftIO)
+import Text.Pandoc.Class (fetchMediaResource, readDataFile, runIO,
+                          runIOorExplode, setUserDataDir)
 import Text.Pandoc.Options (ReaderOptions(readerExtensions))
 import Text.Pandoc.Lua.StackInstances ()
 import Text.Pandoc.Readers (Reader (..), getReader)
+import Text.Pandoc.MIME (MimeType)
+
+import qualified Foreign.Lua as Lua
+import qualified Data.ByteString.Lazy as BL
+import qualified Text.Pandoc.MediaBag as MB
 
 -- | Push the "pandoc" on the lua stack.
 pushPandocModule :: Maybe FilePath -> Lua ()
 pushPandocModule datadir = do
   script <- liftIO (pandocModuleScript datadir)
-  status <- loadstring script
-  unless (status /= OK) $ call 0 1
-  push "__read"
-  pushHaskellFunction readDoc
-  rawset (-3)
+  status <- Lua.loadstring script
+  unless (status /= Lua.OK) $ Lua.call 0 1
+  Lua.push "__read"
+  Lua.pushHaskellFunction readDoc
+  Lua.rawset (-3)
 
 -- | Get the string representation of the pandoc module
 pandocModuleScript :: Maybe FilePath -> IO String
@@ -56,14 +69,98 @@ pandocModuleScript datadir = unpack <$>
 readDoc :: String -> String -> Lua NumResults
 readDoc formatSpec content = do
   case getReader formatSpec of
-    Left  s      -> push s -- Unknown reader
+    Left  s      -> Lua.push s -- Unknown reader
     Right (reader, es) ->
       case reader of
         TextReader r -> do
           res <- liftIO $ runIO $ r def{ readerExtensions = es } (pack content)
           case res of
-            Left s   -> push $ show s -- error while reading
-            Right pd -> push pd       -- success, push Pandoc
-        _  -> push "Only string formats are supported at the moment."
+            Left s   -> Lua.push $ show s -- error while reading
+            Right pd -> Lua.push pd       -- success, push Pandoc
+        _  -> Lua.push "Only string formats are supported at the moment."
+  return 1
+
+--
+-- MediaBag submodule
+--
+pushMediaBagModule :: IORef MB.MediaBag -> Lua ()
+pushMediaBagModule mediaBagRef = do
+  Lua.newtable
+  addFunction "insert" (insertMediaFn mediaBagRef)
+  addFunction "lookup" (lookupMediaFn mediaBagRef)
+  addFunction "list" (mediaDirectoryFn mediaBagRef)
+  addFunction "fetch" (insertResource mediaBagRef)
+  return ()
+ where
+  addFunction name fn = do
+    Lua.push name
+    Lua.pushHaskellFunction fn
+    Lua.rawset (-3)
+
+insertMediaFn :: IORef MB.MediaBag
+              -> FilePath
+              -> OrNil MimeType
+              -> BL.ByteString
+              -> Lua NumResults
+insertMediaFn mbRef fp nilOrMime contents = do
+  liftIO . modifyIORef' mbRef $ MB.insertMedia fp (toMaybe nilOrMime) contents
+  return 0
+
+lookupMediaFn :: IORef MB.MediaBag
+              -> FilePath
+              -> Lua NumResults
+lookupMediaFn mbRef fp = do
+  res <- MB.lookupMedia fp <$> liftIO (readIORef mbRef)
+  case res of
+    Nothing -> Lua.pushnil *> return 1
+    Just (mimeType, contents) -> do
+      Lua.push mimeType
+      Lua.push contents
+      return 2
+
+mediaDirectoryFn :: IORef MB.MediaBag
+                 -> Lua NumResults
+mediaDirectoryFn mbRef = do
+  dirContents <- MB.mediaDirectory <$> liftIO (readIORef mbRef)
+  Lua.newtable
+  zipWithM_ addEntry [1..] dirContents
   return 1
+ where
+  addEntry :: Int -> (FilePath, MimeType, Int) -> Lua ()
+  addEntry idx (fp, mimeType, contentLength) = do
+    Lua.newtable
+    Lua.push "path" *> Lua.push fp *> Lua.rawset (-3)
+    Lua.push "type" *> Lua.push mimeType *> Lua.rawset (-3)
+    Lua.push "length" *> Lua.push contentLength *> Lua.rawset (-3)
+    Lua.rawseti (-2) idx
+
+insertResource :: IORef MB.MediaBag
+               -> String
+               -> OrNil String
+               -> Lua NumResults
+insertResource mbRef src sourceUrlOrNil = do
+  (fp, mimeType, bs) <- liftIO . runIOorExplode $
+                        fetchMediaResource (toMaybe sourceUrlOrNil) src
+  liftIO $ print (fp, mimeType)
+  insertMediaFn mbRef fp (OrNil mimeType) bs
+
+--
+-- Helper types and orphan instances
+--
+
+newtype OrNil a = OrNil { toMaybe :: Maybe a }
+
+instance FromLuaStack a => FromLuaStack (OrNil a) where
+  peek idx = do
+    noValue <- Lua.isnil idx
+    if noValue
+      then return (OrNil Nothing)
+      else OrNil . Just <$> Lua.peek idx
+
+#if !MIN_VERSION_hslua(0,9,0)
+instance ToLuaStack BL.ByteString where
+  push = Lua.push . BL.toStrict
 
+instance FromLuaStack BL.ByteString where
+  peek = fmap BL.fromStrict . Lua.peek
+#endif