{-# LANGUAGE MultiWayIf #-}

module Text.Pandoc.Filter.Pyplot (
        makePlot
      , makePlot'
      , PandocPyplotError(..)
      , showError
    ) where

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

import           Text.Pandoc.Definition

import           Text.Pandoc.Filter.Scripting

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

                       | InvalidTargetError FilePath    -- ^ Invalid figure path

                       | 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)

    , caption :: String     -- ^ Figure caption (optional)

    }

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

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

-- | Determine inclusion specifications from Block attributes.

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

parseFigureSpec :: M.Map String String -> Maybe FigureSpec
parseFigureSpec attrs = createInclusion <$> M.lookup targetKey attrs
    where
        defaultAltText = "Figure generated by pandoc-pyplot"
        defaultCaption = mempty
        createInclusion fname = FigureSpec
            { target  = fname
            , alt     = M.findWithDefault defaultAltText altTextKey attrs
            , caption = M.findWithDefault defaultCaption captionKey attrs
            }

-- | Format the script source based on figure spec.

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

-- | 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' cb @ (CodeBlock (id', cls, attrs) scriptSource) =
    case parseFigureSpec (M.fromList attrs) of
        -- Could not parse - leave code block unchanged

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

        Just spec -> do
            let figurePath = target spec

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

                script <- addPlotCapture figurePath scriptSource
                result <- runTempPythonScript script

                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 = replaceExtension figurePath ".txt"
                        writeFile sourcePath $ formatScriptSource spec scriptSource

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

                        let inclusionKeys = [ targetKey, altTextKey, captionKey ]
                            filteredAttrs = filter (\(k,_) -> k `notElem` inclusionKeys) attrs
                            image         = Image (id', cls, filteredAttrs) [Str $ alt spec] (figurePath, "")
                            srcTarget     = (sourcePath, "Click on this figure to see the source code")

                        -- TODO: use FigureSpec caption

                        -- We make the figure be a link to the source code

                        return $ Right $ Para [
                            Link nullAttr [image] srcTarget
                            ]

makePlot' x = return $ Right x

-- | 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 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