content: {-# LANGUAGE LambdaCase          #-}
  {-# LANGUAGE OverloadedStrings   #-}
  {- |
     Module      : Text.Pandoc.Lua.Writer.Scaffolding
     Copyright   : © 2022-2024 Albert Krewinkel
     License     : GPL-2.0-or-later
     Maintainer  : Albert Krewinkel <albert+pandoc@tarleb.com>
  
  Conversion of Pandoc documents using a custom Lua writer.
  -}
  module Text.Pandoc.Lua.Writer.Scaffolding
    ( pushWriterScaffolding
    ) where
  
  import Control.Monad ((<$!>), void)
  import Data.ByteString (ByteString)
  import Data.Data (dataTypeConstrs, dataTypeOf, showConstr, toConstr)
  import Data.Default (def)
  import Data.List (intersperse)
  import Data.Maybe (fromMaybe)
  import Data.Text (Text)
  import Data.String (IsString (fromString))
  import HsLua
  import HsLua.Module.DocLayout (peekDoc, pushDoc)
  import Text.DocLayout (Doc, blankline, render)
  import Text.DocTemplates (Context)
  import Text.Pandoc.Definition
  import Text.Pandoc.Error (PandocError (..))
  import Text.Pandoc.Options (WriterOptions (..), WrapOption(..))
  import Text.Pandoc.Lua.PandocLua ()
  import Text.Pandoc.Lua.Marshal.AST
  import Text.Pandoc.Lua.Marshal.Context (peekContext)
  import Text.Pandoc.Lua.Marshal.WriterOptions ( peekWriterOptions
                                               , pushWriterOptions)
  import Text.Pandoc.Templates (renderTemplate)
  import Text.Pandoc.Writers.Shared (metaToContext, setField)
  import qualified Data.Text as T
  import qualified Text.Pandoc.UTF8 as UTF8
  
  -- | Convert Pandoc to custom markup.
  pushWriterScaffolding :: LuaE PandocError NumResults
  pushWriterScaffolding = do
    newtable
      *> pushWriterMT *> setmetatable (nth 2)
    writer <- toWriterTable top
    addField "Blocks"  $ pushDocumentedFunction (blocksFn writer)
    addField "Inlines" $ pushDocumentedFunction (inlinesFn writer)
    addField "Block"   $ newtable *> pushBlockMT  writer *> setmetatable (nth 2)
    addField "Inline"  $ newtable *> pushInlineMT writer *> setmetatable (nth 2)
    addField "Pandoc"  $ pushDocumentedFunction $ lambda
      ### (\(Pandoc _ blks) -> do
              pushWriterTable writer
              getfield' top "Blocks"
              pushBlocks blks
              callTrace 1 1
              pure (NumResults 1))
      <#> parameter peekPandoc "Pandoc" "doc" ""
      =?> "rendered doc"
    freeWriter writer
    return 1
   where
    blocksFn w = lambda
      ### (\blocks msep -> blockListToCustom w msep blocks)
      <#> parameter peekBlocks "Blocks" "blocks" ""
      <#> opt (parameter peekDocFuzzy "Doc" "sep" "")
      =#> functionResult pushDoc "Doc" ""
    inlinesFn w = lambda
      ### inlineListToCustom w
      <#> parameter peekInlines "Inlines" "inlines" ""
      =#> functionResult pushDoc "Doc" ""
    pushBlockMT writer = do
      newtable
      addField "__call" $ pushDocumentedFunction $ lambda
        ### blockToCustom
        <#> parameter peekWriter "table" "writer" ""
        <#> parameter peekBlockFuzzy "Block" "block" ""
        =#> functionResult pushDoc "Doc" "rendered blocks"
      addField "__index" $
        -- lookup missing fields in the main Writer table
        pushWriterTable writer
    pushInlineMT writer = do
      newtable
      addField "__call" $ pushDocumentedFunction $ lambda
        ### inlineToCustom
        <#> parameter peekWriter "table" "writer" ""
        <#> parameter peekInlineFuzzy "Inline" "inline" ""
        =#> functionResult pushDoc "Doc" "rendered inline"
      addField "__index" $ do
        -- lookup missing fields in the main Writer table
        pushWriterTable writer
  
  pushWriterMT :: LuaE PandocError ()
  pushWriterMT = do
    newtable
    addField "__call" $ pushDocumentedFunction $ lambda
      ### (\writer doc mopts -> runWriter writer doc mopts)
      <#> parameter peekWriter "table" "writer" ""
      <#> parameter peekPandoc "Pandoc" "doc" ""
      <#> opt (parameter peekWriterOptions "WriterOptions" "opts" "")
      =#> functionResult pushText "string" "rendered document"
    addField "__index" . pushDocumentedFunction $ lambda
      ### (\_writer key -> handleMissingField key)
      <#> parameter pure "table"  "writer" ""
      <#> parameter (liftLua . tostring') "string" "key" ""
      =#> functionResult (const pushnil) "string" ""
  
  
  addField :: LuaError e => Name -> LuaE e a -> LuaE e ()
  addField name action = do
    pushName name
    action
    rawset (nth 3)
  
  getfield' :: LuaError e => StackIndex -> Name -> LuaE e HsLua.Type
  getfield' idx name = do
    aidx <- absindex idx
    pushName name
    rawget aidx >>= \case
      TypeNil -> pop 1 *> getfield aidx name
      ty      -> pure ty
  
  -- | A writer table is just an absolute stack index.
  newtype WriterTable = WriterTable Reference
  
  toWriterTable :: LuaError e => StackIndex -> LuaE e WriterTable
  toWriterTable idx = WriterTable <$!> do
    pushvalue idx
    ref registryindex
  
  peekWriter :: LuaError e => Peeker e WriterTable
  peekWriter = liftLua . toWriterTable
  
  pushWriterTable :: LuaError e => Pusher e WriterTable
  pushWriterTable (WriterTable wref) = void $ getref registryindex wref
  
  writerOptionsField :: Name
  writerOptionsField = "Pandoc Writer WriterOptions"
  
  freeWriter :: WriterTable -> LuaE e ()
  freeWriter (WriterTable wref) = unref registryindex wref
  
  pushOpts :: LuaE PandocError ()
  pushOpts = void $ getfield' registryindex writerOptionsField
  
  runWriter :: WriterTable -> Pandoc -> Maybe WriterOptions
            -> LuaE PandocError Text
  runWriter writer doc@(Pandoc meta _blks) mopts = do
    let opts = fromMaybe def mopts
    pushWriterOptions opts *>
      setfield registryindex writerOptionsField
  
    (body, mcontext) <- runPeek (pandocToCustom writer doc) >>= force . \case
      Failure msg contexts -> Failure (cleanupTrace msg) contexts
      s -> s
  
    -- convert metavalues to a template context (variables)
    defaultContext <- metaToContext opts
                      (blockListToCustom writer Nothing)
                      (inlineListToCustom writer)
                      meta
    let context = setField "body" body
                $ fromMaybe defaultContext mcontext
  
    let colwidth = if writerWrapText opts == WrapAuto
                      then Just $ writerColumns opts
                      else Nothing
  
    return $ render colwidth $
      case writerTemplate opts of
         Nothing  -> body
         Just tpl -> renderTemplate tpl context
  
  -- | Keep exactly one traceback and clean it up. This wouldn't be
  -- necessary if the @pcallTrace@ function would do nothing whenever the
  -- error already included a trace, but that would require some bigger
  -- changes; removing the additional traces in this post-process step is
  -- much easier (for now).
  cleanupTrace :: ByteString -> ByteString
  cleanupTrace msg = UTF8.fromText . T.intercalate "\n" $
    let tmsg = T.lines $ UTF8.toText msg
        traceStart = (== "stack traceback:")
    in case break traceStart tmsg of
          (x, t:traces) -> (x <>) . (t:) $
                           let (firstTrace, rest) = break traceStart traces
                               isPeekContext = ("\twhile " `T.isPrefixOf`)
                               isUnknownCFn = (== "\t[C]: in ?")
                           in filter (not . isUnknownCFn) firstTrace <>
                              filter isPeekContext rest
          _ -> tmsg
  
  -- | Pushes the field in the writer table.
  getWriterField :: LuaError e
                 => WriterTable -> Name -> LuaE e HsLua.Type
  getWriterField writer name = do
    pushWriterTable writer
    getfield' top name <* remove (nth 2)
  
  -- | Looks up @Writer.subtable.field@; tries @Writer.field@ as a fallback if the
  -- subtable field is @nil@.
  getNestedWriterField :: LuaError e
                       => WriterTable -> Name -> Name -> LuaE e HsLua.Type
  getNestedWriterField writer subtable field = do
    pushWriterTable writer
    getfield' top subtable >>= \case
      TypeNil -> TypeNil <$ remove (nth 2) -- remove Writer table
      _       -> getfield' top field
                 -- remove Writer and subtable
                 <* remove (nth 3) <* remove (nth 2)
  
  pandocToCustom :: WriterTable -> Pandoc
                 -> Peek PandocError (Doc Text, Maybe (Context Text))
  pandocToCustom writer doc = withContext "rendering Pandoc" $ do
    callStatus <- liftLua $ do
      getWriterField writer "Pandoc"
      pushPandoc doc
      pushOpts
      pcallTrace 2 2
    case callStatus of
      OK -> ((,) <$> peekDocFuzzy (nth 2) <*> orNil peekContext top)
            `lastly` pop 2
      _  -> failPeek =<< liftLua (tostring' top)
  
  blockToCustom :: WriterTable -> Block -> LuaE PandocError (Doc Text)
  blockToCustom writer blk = forcePeek $ renderBlock writer blk
  
  renderBlock :: WriterTable -> Block -> Peek PandocError (Doc Text)
  renderBlock writer blk = do
    let constrName = fromString . showConstr . toConstr $ blk
    withContext ("rendering Block `" <> constrName <> "`") $
      liftLua (getNestedWriterField writer "Block" constrName) >>= \case
        TypeNil -> failPeek =<< typeMismatchMessage "function or Doc" top
        _       -> callOrDoc (pushBlock blk)
  
  inlineToCustom :: WriterTable -> Inline -> LuaE PandocError (Doc Text)
  inlineToCustom writer inln = forcePeek $ renderInline writer inln
  
  renderInline :: WriterTable -> Inline -> Peek PandocError (Doc Text)
  renderInline writer inln = do
    let constrName = fromString . showConstr . toConstr $ inln
    withContext ("rendering Inline `" <> constrName <> "`") $ do
      liftLua (getNestedWriterField writer "Inline" constrName) >>= \case
        TypeNil -> failPeek =<< typeMismatchMessage "function or Doc" top
        _       -> callOrDoc (pushInline inln)
  
  -- | If the value at the top of the stack can be called as a function,
  -- then push the element and writer options to the stack and call it;
  -- otherwise treat it as a plain Doc value
  callOrDoc :: LuaE PandocError ()
            -> Peek PandocError (Doc Text)
  callOrDoc pushElement = do
    liftLua (ltype top) >>= \case
      TypeFunction -> peekCall
      _            ->
        liftLua (getmetafield top "__call") >>= \case
          TypeNil -> peekDocFuzzy top
          _       -> liftLua (pop 1) *> peekCall
   where
     peekCall :: Peek PandocError (Doc Text)
     peekCall =
       liftLua (pushElement *> pushOpts *> pcallTrace 2 1) >>= \case
         OK -> peekDocFuzzy top
         _  -> failPeek =<< liftLua (tostring' top)
  
  blockListToCustom :: WriterTable -> Maybe (Doc Text) -> [Block]
                    -> LuaE PandocError (Doc Text)
  blockListToCustom writer msep blocks = forcePeek $
    renderBlockList writer msep blocks
  
  inlineListToCustom :: WriterTable -> [Inline] -> LuaE PandocError (Doc Text)
  inlineListToCustom writer inlines = forcePeek $
    renderInlineList writer inlines
  
  renderBlockList :: WriterTable -> Maybe (Doc Text) -> [Block]
                  -> Peek PandocError (Doc Text)
  renderBlockList writer msep blocks = withContext "rendering Blocks" $ do
    let addSeps = intersperse $ fromMaybe blankline msep
    mconcat . addSeps <$> mapM (renderBlock writer) blocks
  
  renderInlineList :: WriterTable -> [Inline] -> Peek PandocError (Doc Text)
  renderInlineList writer inlines = withContext "rendering Inlines" $ do
    mconcat <$> mapM (renderInline writer) inlines
  
  orNil :: Peeker e a -> Peeker e (Maybe a)
  orNil p idx = liftLua (ltype idx) >>= \case
    TypeNil  -> pure Nothing
    TypeNone -> pure Nothing
    _        -> Just <$> p idx
  
  peekDocFuzzy :: LuaError e => Peeker e (Doc Text)
  peekDocFuzzy idx = liftLua (ltype idx) >>= \case
    TypeTable -> mconcat <$!> peekList peekDoc idx
    _         -> peekDoc idx
  
  handleMissingField :: LuaError e => ByteString -> LuaE e ()
  handleMissingField key' =
    let key = UTF8.toString key'
        blockNames  = map (fromString . show) . dataTypeConstrs . dataTypeOf
                        $ HorizontalRule
        inlineNames = map (fromString . show) . dataTypeConstrs . dataTypeOf
                        $ Space
        mtypeName = case () of
         _ | key `elem` blockNames  -> Just "Block"
         _ | key `elem` inlineNames -> Just "Inline"
         _                          -> Nothing
    in case mtypeName of
         Just typeName  -> failLua $
                           "No render function for " <> typeName <> " value " <>
                           "'" <> key <> "';\ndefine a function `Writer." <>
                           typeName <> "." <> key <> "` that returns " <>
                           "a string or Doc."
         _ -> pure ()
  