{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wno-name-shadowing #-}

{- | Hakyll extension for rendering Coq code blocks using Alectryon.

  See the README for instructions on how to use this module.

  Import this module qualified:

  > import qualified Hakyll.Alectryon as Alectryon

  The main two functions are:

  - 'tryTransform': transform Pandoc documents, ignoring those where
    the @alectryon@ field is not set.
  - 'hakyllWith': extend the Hakyll CLI with an option to enable
    processing code blocks with Alectryon. This is disabled by default,
    assuming that the output is already cached.

-}
module Hakyll.Alectryon
  ( -- * Transformations
    tryTransform
  , tryTransform_
  , transform

    -- ** Alectryon and pygments
  , transformAlectryon
  , transformPygments

    -- * Options
  , Options(..)
  , defaultOptions
  , enabledOptions
  , hakyll
  , hakyllWith
  , optionParser

    -- ** Extend Hakyll with custom parsers
  , execParser
  , parserInfo
  , hakyllWithArgsParser

    -- * Low-level interface
  , updateAlectryonBlocks
  , updateBlocks
  , onCoqBlocks
  , onAlectryonBlocks
  , onBlocks
  , doRunAlectryon
  , doRunPygments
  , name

    -- ** Call Alectryon
  , splitAlectryonHTML
  , callAlectryon

    -- ** Call Pygments
  , splitPygmentsHTML
  , callPygments
  ) where

import Control.Applicative (liftA2, (<|>))
import Control.Monad ((>=>))
import Control.Monad.State (State, runState, get, put)
import Data.Foldable (for_)
import qualified Data.Aeson as Aeson
import qualified Data.ByteString.Lazy as LBS (toStrict)
import Data.Text (Text)
import qualified Data.Text as Text
import qualified Data.Text.IO as Text (readFile, writeFile)
import qualified Data.Text.Encoding as Text (decodeUtf8)
import qualified Options.Applicative as OA
import System.FilePath ((</>))
import System.Process (readProcess)

import Text.Pandoc (Block(CodeBlock, RawBlock), Pandoc(..))
import Text.Pandoc.Walk (query, walkM)
import Hakyll.Core.File (TmpFile(..), newTmpFile)
import Hakyll
  ( Rules, Item, Compiler, getMetadata, itemIdentifier
  , withItemBody, unsafeCompiler, lookupString, lookupStringList)
import qualified Hakyll

pygmentsBatch :: String
pygmentsBatch = unlines
  [ "import sys, json"
  , "from pygments import highlight"
  , "from pygments.lexers import CoqLexer"
  , "from pygments.formatters import HtmlFormatter"
  , "code = json.load(sys.stdin)"
  , "for line in code:"
  , "    print(highlight(line, CoqLexer(), HtmlFormatter()))"
  , "    print(\"<!-- pygments-batch-end -->\")"
  ]

-- | If the @alectryon@ flag is set, 'transform'.
--
-- Do nothing if @alectryon@ flag is not set.
tryTransform :: Options -> Item Pandoc -> Compiler (Item Pandoc)
tryTransform opt idoc = do
  m <- getMetadata (itemIdentifier idoc)
  case lookupStringList "alectryon" m <|> enable of
    Just args -> withItemBody (transform opt (lookupString "alectryon-cache" m) args) idoc
    Nothing -> pure idoc
  where
    enable = if enableAll opt then Just [] else Nothing

-- | If the @alectryon@ flag is set, 'transform'.
--
-- Do nothing if @alectryon@ flag is not set.
--
-- This is 'tryTransform' with all options enabled.
tryTransform_ :: Item Pandoc -> Compiler (Item Pandoc)
tryTransform_ = tryTransform enabledOptions

-- | Convert @alectryon@ and @coq@ code blocks to HTML using Alectryon and
-- pygments respectively.
transform ::
  Options ->
  Maybe FilePath {- ^ Cache directory location -} ->
  [String] {- ^ Alectryon options (e.g., Coq @-Q@, @-R@) -} ->
  Pandoc -> Compiler Pandoc
transform opt cache args =
  transformAlectryon opt ((</> "alectryon.html") <$> cache) args
    >=> transformPygments opt ((</> "pygments.html") <$> cache)

-- | Convert @alectryon@ code blocks to HTML using Alectryon.
--
-- > transform opt cache args :: Pandoc -> Compiler Pandoc
transformAlectryon ::
  Options ->
  Maybe FilePath {- ^ Cache file location -} ->
  [String] {- ^ Alectryon options (e.g., Coq @-Q@, @-R@) -} ->
  Pandoc -> Compiler Pandoc
transformAlectryon opt cache args doc = do
    snips <- doRunAlectryon opt cache args blocks
    pure (updateAlectryonBlocks snips doc)
  where
    blocks = query getBlock doc
    getBlock = onAlectryonBlocks [] (: [])

-- | Convert @coq@ code blocks to HTML using pygments.
transformPygments ::
  Options ->
  Maybe FilePath {- ^ Cache file location -} ->
  Pandoc -> Compiler Pandoc
transformPygments opt cache doc = do
    snips <- doRunPygments opt cache blocks
    pure (updateBlocks coqName snips doc)
  where
    blocks = query getBlock doc
    getBlock = onCoqBlocks [] (: [])

-- | @\"alectryon\"@
name :: Text
name = "alectryon"

-- | @\"coq\"@
coqName :: Text
coqName = "coq"

nextBlock :: State [Text] Block
nextBlock = do
  xs <- get
  case xs of
    [] -> error "No blocks left"
    x : xs -> do
      put xs
      pure (RawBlock "html" x)

updateAlectryonBlocks :: [Text] -> Pandoc -> Pandoc
updateAlectryonBlocks = updateBlocks name

updateBlocks :: Text -> [Text] -> Pandoc -> Pandoc
updateBlocks name bs = fst . flip runState bs . walkM setBlock
  where
    setBlock b = onBlocks name (pure b) (\_ -> nextBlock) b

onCoqBlocks :: b -> (Text -> b) -> Block -> b
onCoqBlocks = onBlocks coqName

onAlectryonBlocks :: b -> (Text -> b) -> Block -> b
onAlectryonBlocks = onBlocks name

onBlocks :: Text -> b -> (Text -> b) -> Block -> b
onBlocks name _ f (CodeBlock (_, cs, _) b) | name `elem` cs = f b
onBlocks _ x _ _ = x

-- | Transform a list of blocks using alectryon.
--
-- If the input list is empty, return the empty list immediately,
-- saving an external call.
doRunAlectryon :: Options -> Maybe FilePath -> [String] -> [Text] -> Compiler [Text]
doRunAlectryon _ _ _ [] = pure []
doRunAlectryon opt cache args blocks | runAlectryon opt = do
  TmpFile snippets <- case cache of
    Just snippets -> pure (TmpFile snippets)
    Nothing -> newTmpFile "snippets.html"
  Hakyll.unsafeCompiler $ do
    callAlectryon args snippets (Text.decodeUtf8 (LBS.toStrict (Aeson.encode blocks)))
    splitAlectryonHTML <$> Text.readFile snippets
doRunAlectryon _opt cache _args _blocks | otherwise =
  case cache of
    Nothing -> error "hakyll-alectryon: cache not set, enable --run-alectryon"
    Just snippets -> Hakyll.unsafeCompiler (splitAlectryonHTML <$> Text.readFile snippets)

-- | Transform a list of blocks using pygments.
--
-- If the input list is empty, return the empty list immediately,
-- saving an external call.
doRunPygments :: Options -> Maybe FilePath -> [Text] -> Compiler [Text]
doRunPygments _ _ [] = pure []
doRunPygments opt cache blocks | runAlectryon opt = do
  Hakyll.unsafeCompiler $ do
    out <- callPygments (Text.decodeUtf8 (LBS.toStrict (Aeson.encode blocks)))
    for_ cache (\snippets -> Text.writeFile snippets out)
    pure (splitPygmentsHTML out)
doRunPygments _opt cache _blocks | otherwise =
  case cache of
    Nothing -> error "hakyll-alectryon: cache not set, enable --run-alectryon"
    Just snippets -> Hakyll.unsafeCompiler (splitPygmentsHTML <$> Text.readFile snippets)

-- | Split Alectryon output.
splitAlectryonHTML :: Text -> [Text]
splitAlectryonHTML = Text.splitOn "<!-- alectryon-block-end -->"

-- | Split pygments output.
splitPygmentsHTML :: Text -> [Text]
splitPygmentsHTML = Text.splitOn "<!-- pygments-batch-end -->"

-- | Call the Alectryon executable.
callAlectryon :: [String] -> FilePath -> Text -> IO ()
callAlectryon args snippets x = do
  _ <- readProcess
    "alectryon"
    (["--frontend", "coq.json", "--backend", "snippets-html", "--stdin-filename", "stdin", "-o", snippets, "-"] ++ args)
    (Text.unpack x)
  pure ()

-- | Call the pygments script.
callPygments :: Text -> IO Text
callPygments x = do
  Text.pack <$> readProcess "python3" ["-c", pygmentsBatch] (Text.unpack x)

-- | Configurations for hakyll-alectryon.
--
-- Avoid using the 'Options' constructor directly, prefer record updates on
-- 'defaultOptions'.
data Options = Options
  { -- | If @True@, actually process code blocks using Alectryon and Pygments.
    -- If @False@, use the cached output, and fail if the cache is not found
    -- (this is to ensure the site is buildable without Alectryon).
    runAlectryon :: Bool

    -- | If @True@, process all posts. If @False@, process only posts whose
    -- metadata contains an @alectryon@ field set to a list.
    --
    -- That list contains options to pass to the Alectryon executable.
    -- It defaults to the empty list if @enableAll = True@ and it is not
    -- present in the metadata of a post.
  , enableAll :: Bool
  }

defaultOptions :: Options
defaultOptions = Options
  { runAlectryon = False
  , enableAll = False
  }

-- | Always run Alectryon on all posts.
enabledOptions :: Options
enabledOptions = Options
  { runAlectryon = True
  , enableAll = True
  }

optionParser :: OA.Parser Options
optionParser = Options
  <$> OA.switch (OA.long "run-alectryon" <> OA.help "Process Alectryon-enabled posts")
  <*> OA.switch (OA.long "enable-all" <> OA.help "Enable Alectryon on all posts")

execParser :: OA.Parser a -> Hakyll.Configuration -> IO (Hakyll.Options, a)
execParser p config = OA.customExecParser (OA.prefs OA.showHelpOnError) (parserInfo p config)

parserInfo :: OA.Parser a -> Hakyll.Configuration -> OA.ParserInfo (Hakyll.Options, a)
parserInfo p config = info { OA.infoParser = liftA2 (,) (OA.infoParser info) p }
  where
    info = Hakyll.defaultParserInfo config

hakyllWithArgsParser :: OA.Parser a -> Hakyll.Configuration -> (a -> Rules ()) -> IO ()
hakyllWithArgsParser oparser config rules = do
  (hopts, opts) <- execParser oparser config
  Hakyll.hakyllWithArgs config hopts (rules opts)

hakyllWith :: Hakyll.Configuration -> (Options -> Rules ()) -> IO ()
hakyllWith = hakyllWithArgsParser optionParser

hakyll :: (Options -> Rules ()) -> IO ()
hakyll = hakyllWith Hakyll.defaultConfiguration