Notebook.hs

module ScriptHs.Notebook where

import Data.Bifunctor (Bifunctor (second))
import Data.Text (Text)
import qualified Data.Text as T
import qualified Data.Text.IO as TIO

import ScriptHs.Markdown (
    CodeOutput (..),
    MimeType (..),
    Segment (..),
    parseMarkdown,
    reassemble,
 )
import ScriptHs.Parser (
    CabalMeta (..),
    Line (..),
    ScriptFile (..),
    mergeMetas,
    parseScript,
 )
import ScriptHs.Run (RunOptions, runScriptCapture)

type IndexedSegments = [(Int, Segment)]
type IndexedBlocks = [(Int, [Line])]

runNotebook :: RunOptions -> FilePath -> Maybe FilePath -> IO ()
runNotebook opts path outputPath = do
    contents <- TIO.readFile path
    outputMd <- processNotebook opts path contents
    case outputPath of
        Nothing -> TIO.putStr outputMd
        Just output -> TIO.writeFile output outputMd

processNotebook :: RunOptions -> FilePath -> Text -> IO Text
processNotebook opts notebookPath contents = do
    let indexedSegments = zip [0 ..] (parseMarkdown contents)
        (metas, indexedCodeBlocks) = parseBlocks indexedSegments
    if null indexedCodeBlocks
        then pure contents
        else executeCodeCells opts notebookPath metas indexedSegments indexedCodeBlocks

executeCodeCells ::
    RunOptions ->
    FilePath ->
    CabalMeta ->
    IndexedSegments ->
    IndexedBlocks ->
    IO Text
executeCodeCells opts notebookPath meta allSegments codeBlocks = do
    let ghciScript = generatedMarkedScript codeBlocks
        sf = ScriptFile{scriptMeta = meta, scriptLines = ghciScript}
    rawOutput <- runScriptCapture opts notebookPath sf
    let outputs = splitByMarkers rawOutput (map fst codeBlocks)
        blocksWithOutput = addOutputToSegments outputs allSegments
    pure $ reassemble blocksWithOutput

addOutputToSegments :: [(Int, Text)] -> IndexedSegments -> [Segment]
addOutputToSegments outputs = map addOutput
  where
    addOutput :: (Int, Segment) -> Segment
    addOutput (i, CodeBlock lang code _) = CodeBlock lang code (fmap (CodeOutput MimePlain) (lookup i outputs))
    addOutput (_, seg) = seg

mkIndexedCodeSegments :: IndexedSegments -> IndexedSegments
mkIndexedCodeSegments segments = [(i, c) | (i, c@(CodeBlock lang _ _)) <- segments, isHaskell lang]

parseBlocks :: IndexedSegments -> (CabalMeta, IndexedBlocks)
parseBlocks blocks = (metas, map (second scriptLines) sfs)
  where
    sfs =
        [(i, parseScript code) | (i, CodeBlock lang code _) <- blocks, isHaskell lang]
    metas = mergeMetas (map (scriptMeta . snd) sfs)

generatedMarkedScript :: IndexedBlocks -> [Line]
generatedMarkedScript = concatMap renderWithMarker
  where
    renderWithMarker :: (Int, [Line]) -> [Line]
    renderWithMarker (idx, ls) =
        ls
            ++ [ Blank
               , HaskellLine ("putStrLn " <> T.pack (show (T.unpack (mkMarker idx))))
               , Blank
               ]

splitByMarkers :: Text -> [Int] -> [(Int, Text)]
splitByMarkers _ [] = []
splitByMarkers remaining (idx : rest) =
    let marker = mkMarker idx
        (before, after) = T.breakOn marker remaining
     in if T.null after
            then [(idx, T.strip before)]
            else
                (idx, T.strip before) : splitByMarkers (T.drop (T.length marker) after) rest

-- A marker we'll print to GHCi output to
-- denote the end of a cell execution block.
mkMarker :: Int -> Text
mkMarker n = "---SCRIPTHS_BLOCK_" <> T.pack (show n) <> "_END---"

isHaskell :: Text -> Bool
isHaskell lang = fenceLanguage lang `elem` ["haskell", "hs"]

isPython :: Text -> Bool
isPython lang = fenceLanguage lang `elem` ["python", "py"]

{- | The base language name of a code-fence info string, lower-cased and
trimmed. Handles both bare tags (@haskell@, @hs@) and Pandoc-style attribute
tags: @{haskell}@, @{.haskell}@, @{.haskell:hs}@, @{.haskell:ghci}@.
-}
fenceLanguage :: Text -> Text
fenceLanguage lang =
    let stripped = T.strip lang
        inner = case T.stripPrefix "{" stripped of
            Just s -> T.takeWhile (/= '}') s
            Nothing -> stripped
        token = case T.words inner of
            (t : _) -> t
            [] -> ""
        base = T.takeWhile (/= ':') (T.dropWhile (== '.') token)
     in T.toLower base