{-# LANGUAGE MultiWayIf #-}
{-# LANGUAGE Unsafe     #-}
{-|
Module      : Text.Pandoc.Filter.Pyplot
Description : Pandoc filter to create Matplotlib figures from code blocks
Copyright   : (c) Laurent P René de Cotret, 2018
License     : MIT
Maintainer  : laurent.decotret@outlook.com
Stability   : stable
Portability : portable

This module defines a Pandoc filter @makePlot@ that can be 
used to walk over a Pandoc document and generate figures from
Python code blocks.
-}
module Text.Pandoc.Filter.Pyplot (
        makePlot
      , makePlot'
      , plotTransform
      , PandocPyplotError(..)
      , showError
    ) where

import           Control.Monad                  ((>=>))
import qualified Data.Map.Strict                as M
import           System.Directory               (doesDirectoryExist)
import           System.FilePath                (isValid, replaceExtension, takeDirectory)

import           Text.Pandoc.Definition
import           Text.Pandoc.Walk               (walkM)

import           Text.Pandoc.Filter.Scripting

-- | Possible errors returned by the filter

data PandocPyplotError = ScriptError Int                -- ^ Running Python script has yielded an error

                       | InvalidTargetError FilePath    -- ^ Invalid figure path

                       | MissingDirectoryError FilePath -- ^ Directory where to save figure does not exist

                       | BlockingCallError              -- ^ Python script contains a block call to 'show()'


-- | Datatype containing all parameters required

-- to run pandoc-pyplot

data FigureSpec = FigureSpec
    { target  :: FilePath     -- ^ filepath where generated figure will be saved

    , alt     :: String       -- ^ Alternate text for the figure (optional)

    , script  :: PythonScript -- ^ Source code for the figure

    , blockAttrs :: Attr      -- ^ Attributes not related to @pandoc-pyplot@ will be propagated

    }

-- | Get the source code for a script including provisions to capture

-- the output.

scriptWithCapture :: FigureSpec -> PythonScript
scriptWithCapture spec = addPlotCapture (target spec) (script spec)

-- | Determine where to save the script source based on plot target

scriptSourcePath :: FigureSpec -> FilePath
scriptSourcePath spec = replaceExtension (target spec) ".txt"

-- | Get the source code for a figure script in a presentable way

presentableScript :: FigureSpec -> PythonScript
presentableScript spec = mconcat [ "# Source code for ", target spec, "\n", script spec ]

-- Keys that pandoc-pyplot will look for in code blocks

targetKey, altTextKey :: String
targetKey  = "plot_target"
altTextKey = "plot_alt"

-- | Determine inclusion specifications from Block attributes.

-- Note that the target key is required, but all other parameters are optional

parseFigureSpec :: Block -> Maybe FigureSpec
parseFigureSpec (CodeBlock (id', cls, attrs) content) =
    createInclusion <$> M.lookup targetKey attrs'
    where
        attrs' = M.fromList attrs
        inclusionKeys = [ targetKey, altTextKey ]
        filteredAttrs = filter (\(k,_) -> k `notElem` inclusionKeys) attrs
        createInclusion fname = FigureSpec
            { target     = fname
            , alt        = M.findWithDefault "Figure generated by pandoc-pyplot" altTextKey attrs'
            , script     = content
            -- Propagate attributes that are not related to pandoc-pyplot

            , blockAttrs = (id', cls, filteredAttrs)
            }
parseFigureSpec _ = Nothing

-- | Main routine to include Matplotlib plots.

-- Code blocks containing the attributes @plot_target@ are considered

-- Python plotting scripts. All other possible blocks are ignored.

-- The source code is also saved in another file, which can be access by 

-- clicking the image

makePlot' :: Block -> IO (Either PandocPyplotError Block)
makePlot' block =
    case parseFigureSpec block of
        -- Could not parse - leave code block unchanged

        Nothing -> return $ Right block
        -- Could parse : run the script and capture output

        Just spec -> do
            let figurePath = target spec
                figureDir = takeDirectory figurePath
                scriptSource = script spec

            -- Check that the directory in which to save the figure exists

            validDirectory <- doesDirectoryExist figureDir

            if | not (isValid figurePath)         -> return $ Left $ InvalidTargetError figurePath
               | not validDirectory               -> return $ Left $ MissingDirectoryError figureDir
               | hasBlockingShowCall scriptSource -> return $ Left $ BlockingCallError
               | otherwise -> do

                -- Running the script happens on the next line

                -- Note that the script is slightly modified to be able to capture the output

                result <- runTempPythonScript (scriptWithCapture spec)

                case result of
                    ScriptFailure code -> return $ Left $ ScriptError code
                    ScriptSuccess -> do
                        -- Save the original script into a separate file

                        -- so it can be inspected

                        -- Note : using a .txt file allows to view source directly

                        --        in the browser, in the case of HTML output

                        let sourcePath = scriptSourcePath spec
                        writeFile sourcePath (presentableScript spec)

                        -- Propagate attributes that are not related to pandoc-pyplot

                        let relevantAttrs = blockAttrs spec
                            srcTarget = Link nullAttr [Str "Source code"] (sourcePath, "")
                            caption'   = [Str $ alt spec, Space, Str "(", srcTarget, Str ")"]
                            -- To render images as figures with captions, the target title

                            -- must be "fig:"

                            -- Janky? yes

                            image     = Image relevantAttrs caption' (figurePath, "fig:")

                        return $ Right $ Para $ [image]

-- | Translate filter error to an error message

showError :: PandocPyplotError -> String
showError (ScriptError exitcode)          = "Script error: plot could not be generated. Exit code " <> (show exitcode)
showError (InvalidTargetError fname)      = "Target filename " <> fname <> " is not valid."
showError (MissingDirectoryError dirname) = "Target directory " <> dirname <> " does not exist."
showError BlockingCallError               = "Script contains a blocking call to show, like 'plt.show()'"

-- | Highest-level function that can be walked over a Pandoc tree.

-- All code blocks that have the 'plot_target' parameter will be considered

-- figures.

makePlot :: Block -> IO Block
makePlot = makePlot' >=> either (fail . showError) return

-- | Walk over an entire Pandoc document, changing appropriate code blocks

-- into figures.

plotTransform :: Pandoc -> IO Pandoc
plotTransform = walkM makePlot