content: -- | @futhark literate@
  --
  -- Also contains various utility definitions used by "Futhark.CLI.Script".
  module Futhark.CLI.Literate
    ( main,
      Options (..),
      initialOptions,
      scriptCommandLineOptions,
      prepareServer,
    )
  where
  
  import Codec.BMP qualified as BMP
  import Control.Monad
  import Control.Monad.Except
  import Control.Monad.State hiding (State)
  import Data.Bifunctor (first, second)
  import Data.Bits
  import Data.ByteString qualified as BS
  import Data.ByteString.Lazy qualified as LBS
  import Data.Char
  import Data.Functor (($>))
  import Data.Int (Int64)
  import Data.List qualified as L
  import Data.Map qualified as M
  import Data.Maybe
  import Data.Set qualified as S
  import Data.Text qualified as T
  import Data.Text.Encoding qualified as T
  import Data.Text.IO qualified as T
  import Data.Text.Read qualified as T
  import Data.Vector.Storable qualified as SVec
  import Data.Vector.Storable.ByteString qualified as SVec
  import Data.Void
  import Data.Word (Word32, Word8)
  import Futhark.Data
  import Futhark.Script
  import Futhark.Server
  import Futhark.Test
  import Futhark.Test.Values
  import Futhark.Util
    ( directoryContents,
      ensureCacheDirectory,
      fancyTerminal,
      hashText,
      nubOrd,
      runProgramWithExitCode,
      showText,
    )
  import Futhark.Util.Options
  import Futhark.Util.Pretty (prettyText, prettyTextOneLine)
  import Futhark.Util.Pretty qualified as PP
  import Futhark.Util.ProgressBar
  import System.Directory
    ( copyFile,
      doesFileExist,
      getCurrentDirectory,
      removePathForcibly,
      setCurrentDirectory,
    )
  import System.Environment (getExecutablePath)
  import System.Exit
  import System.FilePath
  import System.IO
  import System.IO.Error (isDoesNotExistError)
  import System.IO.Temp (withSystemTempDirectory, withSystemTempFile)
  import Text.Megaparsec hiding (State, failure, token)
  import Text.Megaparsec.Char
  import Text.Printf
  
  newtype ImgParams = ImgParams
    { imgFile :: Maybe FilePath
    }
    deriving (Show)
  
  defaultImgParams :: ImgParams
  defaultImgParams =
    ImgParams {imgFile = Nothing}
  
  data VideoParams = VideoParams
    { videoFPS :: Maybe Int,
      videoLoop :: Maybe Bool,
      videoAutoplay :: Maybe Bool,
      videoFormat :: Maybe T.Text,
      videoFile :: Maybe FilePath
    }
    deriving (Show)
  
  defaultVideoParams :: VideoParams
  defaultVideoParams =
    VideoParams
      { videoFPS = Nothing,
        videoLoop = Nothing,
        videoAutoplay = Nothing,
        videoFormat = Nothing,
        videoFile = Nothing
      }
  
  data AudioParams = AudioParams
    { audioSamplingFrequency :: Maybe Int,
      audioCodec :: Maybe T.Text
    }
    deriving (Show)
  
  defaultAudioParams :: AudioParams
  defaultAudioParams =
    AudioParams
      { audioSamplingFrequency = Nothing,
        audioCodec = Nothing
      }
  
  data Directive
    = DirectiveRes Exp
    | DirectiveBrief Directive
    | DirectiveCovert Directive
    | DirectiveImg Exp ImgParams
    | DirectivePlot Exp (Maybe (Int, Int))
    | DirectiveGnuplot Exp T.Text
    | DirectiveVideo Exp VideoParams
    | DirectiveAudio Exp AudioParams
    deriving (Show)
  
  varsInDirective :: Directive -> S.Set EntryName
  varsInDirective (DirectiveRes e) = varsInExp e
  varsInDirective (DirectiveBrief d) = varsInDirective d
  varsInDirective (DirectiveCovert d) = varsInDirective d
  varsInDirective (DirectiveImg e _) = varsInExp e
  varsInDirective (DirectivePlot e _) = varsInExp e
  varsInDirective (DirectiveGnuplot e _) = varsInExp e
  varsInDirective (DirectiveVideo e _) = varsInExp e
  varsInDirective (DirectiveAudio e _) = varsInExp e
  
  pprDirective :: Bool -> Directive -> PP.Doc a
  pprDirective _ (DirectiveRes e) =
    "> " <> PP.align (PP.pretty e)
  pprDirective _ (DirectiveBrief f) =
    pprDirective False f
  pprDirective _ (DirectiveCovert f) =
    pprDirective False f
  pprDirective _ (DirectiveImg e params) =
    ("> :img " <> PP.align (PP.pretty e))
      <> if null params' then mempty else ";" <> PP.hardline <> PP.stack params'
    where
      params' = catMaybes [p "file" imgFile PP.pretty]
      p s f pretty = do
        x <- f params
        Just $ s <> ": " <> pretty x
  pprDirective True (DirectivePlot e (Just (h, w))) =
    PP.stack
      [ "> :plot2d " <> PP.pretty e <> ";",
        "size: (" <> PP.pretty w <> "," <> PP.pretty h <> ")"
      ]
  pprDirective _ (DirectivePlot e _) =
    "> :plot2d " <> PP.align (PP.pretty e)
  pprDirective True (DirectiveGnuplot e script) =
    PP.stack $
      "> :gnuplot " <> PP.align (PP.pretty e) <> ";"
        : map PP.pretty (T.lines script)
  pprDirective False (DirectiveGnuplot e _) =
    "> :gnuplot " <> PP.align (PP.pretty e)
  pprDirective False (DirectiveVideo e _) =
    "> :video " <> PP.align (PP.pretty e)
  pprDirective True (DirectiveVideo e params) =
    ("> :video " <> PP.pretty e)
      <> if null params' then mempty else ";" <> PP.hardline <> PP.stack params'
    where
      params' =
        catMaybes
          [ p "fps" videoFPS PP.pretty,
            p "loop" videoLoop ppBool,
            p "autoplay" videoAutoplay ppBool,
            p "format" videoFormat PP.pretty,
            p "file" videoFile PP.pretty
          ]
      ppBool b = if b then "true" else "false"
      p s f pretty = do
        x <- f params
        Just $ s <> ": " <> pretty x
  pprDirective _ (DirectiveAudio e params) =
    ("> :audio " <> PP.pretty e)
      <> if null params' then mempty else ";" <> PP.hardline <> PP.stack params'
    where
      params' =
        catMaybes
          [ p "sampling_frequency" audioSamplingFrequency PP.pretty,
            p "codec" audioCodec PP.pretty
          ]
      p s f pretty = do
        x <- f params
        Just $ s <> ": " <> pretty x
  
  instance PP.Pretty Directive where
    pretty = pprDirective True
  
  data Block
    = BlockCode T.Text
    | BlockComment T.Text
    | BlockDirective Directive T.Text
    deriving (Show)
  
  varsInScripts :: [Block] -> S.Set EntryName
  varsInScripts = foldMap varsInBlock
    where
      varsInBlock (BlockDirective d _) = varsInDirective d
      varsInBlock BlockCode {} = mempty
      varsInBlock BlockComment {} = mempty
  
  type Parser = Parsec Void T.Text
  
  postlexeme :: Parser ()
  postlexeme = void $ hspace *> optional (try $ eol *> "--" *> postlexeme)
  
  lexeme :: Parser a -> Parser a
  lexeme p = p <* postlexeme
  
  token :: T.Text -> Parser ()
  token = void . try . lexeme . string
  
  parseInt :: Parser Int
  parseInt = lexeme $ read <$> some (satisfy isDigit)
  
  restOfLine :: Parser T.Text
  restOfLine = takeWhileP Nothing (/= '\n') <* (void eol <|> eof)
  
  parseBlockComment :: Parser T.Text
  parseBlockComment = T.unlines <$> some line
    where
      line = "--" *> optional " " *> restOfLine
  
  parseTestBlock :: Parser T.Text
  parseTestBlock =
    T.unlines <$> ((:) <$> header <*> remainder)
    where
      header = "-- ==" <* eol
      remainder = map ("-- " <>) . T.lines <$> parseBlockComment
  
  parseBlockCode :: Parser T.Text
  parseBlockCode = T.unlines . noblanks <$> some line
    where
      noblanks = reverse . dropWhile T.null . reverse . dropWhile T.null
      line = try (notFollowedBy "--") *> notFollowedBy eof *> restOfLine
  
  parsePlotParams :: Parser (Maybe (Int, Int))
  parsePlotParams =
    optional $
      ";"
        *> hspace
        *> eol
        *> token "-- size:"
        *> token "("
        *> ((,) <$> parseInt <* token "," <*> parseInt)
        <* token ")"
  
  withPredicate :: (a -> Bool) -> String -> Parser a -> Parser a
  withPredicate f msg p = do
    r <- lookAhead p
    if f r then p else fail msg
  
  parseFilePath :: Parser FilePath
  parseFilePath =
    withPredicate ok "filename must not have directory component" p
    where
      p = T.unpack <$> lexeme (takeWhileP Nothing (not . isSpace))
      ok f = takeFileName f == f
  
  parseImgParams :: Parser ImgParams
  parseImgParams =
    fmap (fromMaybe defaultImgParams) $
      optional $
        ";" *> hspace *> eol *> "-- " *> parseParams defaultImgParams
    where
      parseParams params =
        choice
          [ choice
              [pFile params]
              >>= parseParams,
            pure params
          ]
      pFile params = do
        token "file:"
        b <- parseFilePath
        pure params {imgFile = Just b}
  
  parseVideoParams :: Parser VideoParams
  parseVideoParams =
    fmap (fromMaybe defaultVideoParams) $
      optional $
        ";" *> hspace *> eol *> "-- " *> parseParams defaultVideoParams
    where
      parseParams params =
        choice
          [ choice
              [pLoop params, pFPS params, pAutoplay params, pFormat params]
              >>= parseParams,
            pure params
          ]
      parseBool = token "true" $> True <|> token "false" $> False
      pLoop params = do
        token "loop:"
        b <- parseBool
        pure params {videoLoop = Just b}
      pFPS params = do
        token "fps:"
        fps <- parseInt
        pure params {videoFPS = Just fps}
      pAutoplay params = do
        token "autoplay:"
        b <- parseBool
        pure params {videoAutoplay = Just b}
      pFormat params = do
        token "format:"
        s <- lexeme $ takeWhileP Nothing (not . isSpace)
        pure params {videoFormat = Just s}
  
  parseAudioParams :: Parser AudioParams
  parseAudioParams =
    fmap (fromMaybe defaultAudioParams) $
      optional $
        ";" *> hspace *> eol *> "-- " *> parseParams defaultAudioParams
    where
      parseParams params =
        choice
          [ choice
              [pSamplingFrequency params, pCodec params]
              >>= parseParams,
            pure params
          ]
      pSamplingFrequency params = do
        token "sampling_frequency:"
        hz <- parseInt
        pure params {audioSamplingFrequency = Just hz}
      pCodec params = do
        token "codec:"
        s <- lexeme $ takeWhileP Nothing (not . isSpace)
        pure params {audioCodec = Just s}
  
  atStartOfLine :: Parser ()
  atStartOfLine = do
    col <- sourceColumn <$> getSourcePos
    when (col /= pos1) empty
  
  afterExp :: Parser ()
  afterExp = choice [atStartOfLine, choice [void eol, eof]]
  
  withParsedSource :: Parser a -> (a -> T.Text -> b) -> Parser b
  withParsedSource p f = do
    s <- getInput
    bef <- getOffset
    x <- p
    aft <- getOffset
    pure $ f x $ T.take (aft - bef) s
  
  stripCommentPrefix :: T.Text -> T.Text
  stripCommentPrefix = T.unlines . map onLine . T.lines
    where
      onLine s
        | "-- " `T.isPrefixOf` s = T.drop 3 s
        | otherwise = T.drop 2 s
  
  parseBlock :: Parser Block
  parseBlock =
    choice
      [ withParsedSource (token "-- >" *> parseDirective) $ \d s ->
          BlockDirective d $ stripCommentPrefix s,
        BlockCode <$> parseTestBlock,
        BlockCode <$> parseBlockCode,
        BlockComment <$> parseBlockComment
      ]
    where
      parseDirective =
        choice
          [ DirectiveRes <$> parseExp postlexeme <* afterExp,
            directiveName "covert"
              $> DirectiveCovert
              <*> parseDirective,
            directiveName "brief"
              $> DirectiveBrief
              <*> parseDirective,
            directiveName "img"
              $> DirectiveImg
              <*> parseExp postlexeme
              <*> parseImgParams
              <* choice [void eol, eof],
            directiveName "plot2d"
              $> DirectivePlot
              <*> parseExp postlexeme
              <*> parsePlotParams
              <* choice [void eol, eof],
            directiveName "gnuplot"
              $> DirectiveGnuplot
              <*> parseExp postlexeme
              <*> (";" *> hspace *> eol *> parseBlockComment),
            (directiveName "video" <|> directiveName "video")
              $> DirectiveVideo
              <*> parseExp postlexeme
              <*> parseVideoParams
              <* eol,
            directiveName "audio"
              $> DirectiveAudio
              <*> parseExp postlexeme
              <*> parseAudioParams
              <* choice [void eol, eof]
          ]
      directiveName s = try $ token (":" <> s)
  
  parseProg :: FilePath -> T.Text -> Either T.Text [Block]
  parseProg fname s =
    either (Left . T.pack . errorBundlePretty) Right $
      parse (many parseBlock <* eof) fname s
  
  parseProgFile :: FilePath -> IO [Block]
  parseProgFile prog = do
    pres <- parseProg prog <$> T.readFile prog
    case pres of
      Left err -> do
        T.hPutStr stderr err
        exitFailure
      Right script ->
        pure script
  
  -- | The collection of file paths (all inside the image directory)
  -- produced during directive execution.
  type Files = S.Set FilePath
  
  newtype State = State {stateFiles :: Files}
  
  newtype ScriptM a = ScriptM (ExceptT T.Text (StateT State IO) a)
    deriving
      ( Functor,
        Applicative,
        Monad,
        MonadError T.Text,
        MonadFail,
        MonadIO,
        MonadState State
      )
  
  runScriptM :: ScriptM a -> IO (Either T.Text a, Files)
  runScriptM (ScriptM m) = second stateFiles <$> runStateT (runExceptT m) s
    where
      s = State mempty
  
  withTempFile :: (FilePath -> ScriptM a) -> ScriptM a
  withTempFile f =
    join . liftIO . withSystemTempFile "futhark-literate" $ \tmpf tmpf_h -> do
      hClose tmpf_h
      (res, files) <- runScriptM (f tmpf)
      pure $ do
        modify $ \s -> s {stateFiles = files <> stateFiles s}
        either throwError pure res
  
  withTempDir :: (FilePath -> ScriptM a) -> ScriptM a
  withTempDir f =
    join . liftIO . withSystemTempDirectory "futhark-literate" $ \dir -> do
      (res, files) <- runScriptM (f dir)
      pure $ do
        modify $ \s -> s {stateFiles = files <> stateFiles s}
        either throwError pure res
  
  greyFloatToImg ::
    (RealFrac a, SVec.Storable a) =>
    SVec.Vector a ->
    SVec.Vector Word32
  greyFloatToImg = SVec.map grey
    where
      grey i =
        let i' = round (i * 255) .&. 0xFF
         in (i' `shiftL` 16) .|. (i' `shiftL` 8) .|. i'
  
  greyByteToImg ::
    (Integral a, SVec.Storable a) =>
    SVec.Vector a ->
    SVec.Vector Word32
  greyByteToImg = SVec.map grey
    where
      grey i =
        (fromIntegral i `shiftL` 16) .|. (fromIntegral i `shiftL` 8) .|. fromIntegral i
  
  -- BMPs are RGBA and bottom-up where we assumes images are top-down
  -- and ARGB.  We fix this up before encoding the BMP.  This is
  -- probably a little slower than it has to be.
  vecToBMP :: Int -> Int -> SVec.Vector Word32 -> LBS.ByteString
  vecToBMP h w = BMP.renderBMP . BMP.packRGBA32ToBMP24 w h . SVec.vectorToByteString . frobVec
    where
      frobVec vec = SVec.generate (h * w * 4) (pix vec)
      pix vec l =
        let (i, j) = (l `div` 4) `divMod` w
            argb = vec SVec.! ((h - 1 - i) * w + j)
            c = (argb `shiftR` (24 - ((l + 1) `mod` 4) * 8)) .&. 0xFF
         in fromIntegral c :: Word8
  
  valueToBMP :: Value -> Maybe LBS.ByteString
  valueToBMP v@(U32Value _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w bytes
  valueToBMP v@(I32Value _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w $ SVec.map fromIntegral bytes
  valueToBMP v@(F32Value _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w $ greyFloatToImg bytes
  valueToBMP v@(U8Value _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w $ greyByteToImg bytes
  valueToBMP v@(F64Value _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w $ greyFloatToImg bytes
  valueToBMP v@(BoolValue _ bytes)
    | [h, w] <- valueShape v =
        Just $ vecToBMP h w $ greyByteToImg $ SVec.map ((*) 255 . fromEnum) bytes
  valueToBMP _ = Nothing
  
  valueToBMPs :: Value -> Maybe [LBS.ByteString]
  valueToBMPs = mapM valueToBMP . valueElems
  
  system ::
    (MonadIO m, MonadError T.Text m) =>
    FilePath ->
    [String] ->
    T.Text ->
    m T.Text
  system prog options input = do
    res <- liftIO $ runProgramWithExitCode prog options $ T.encodeUtf8 input
    case res of
      Left err ->
        throwError $ prog' <> " failed: " <> showText err
      Right (ExitSuccess, stdout_t, _) ->
        pure $ T.pack stdout_t
      Right (ExitFailure code', _, stderr_t) ->
        throwError $
          prog'
            <> " failed with exit code "
            <> showText code'
            <> " and stderr:\n"
            <> T.pack stderr_t
    where
      prog' = "'" <> T.pack prog <> "'"
  
  formatDataForGnuplot :: [Value] -> T.Text
  formatDataForGnuplot = T.unlines . map line . L.transpose . map valueElems
    where
      line = T.unwords . map prettyText
  
  imgBlock :: FilePath -> T.Text
  imgBlock f = "![](" <> T.pack f <> ")\n"
  
  videoBlock :: VideoParams -> FilePath -> T.Text
  videoBlock opts f = "![](" <> T.pack f <> ")" <> opts' <> "\n"
    where
      opts'
        | all T.null [loop, autoplay] =
            mempty
        | otherwise =
            "{" <> T.unwords [loop, autoplay] <> "}"
      boolOpt s prop
        | Just b <- prop opts =
            if b then s <> "=\"true\"" else s <> "=\"false\""
        | otherwise =
            mempty
      loop = boolOpt "loop" videoLoop
      autoplay = boolOpt "autoplay" videoAutoplay
  
  plottable :: CompoundValue -> Maybe [Value]
  plottable (ValueTuple vs) = do
    (vs', ns') <- mapAndUnzipM inspect vs
    guard $ length (nubOrd ns') == 1
    Just vs'
    where
      inspect (ValueAtom v)
        | [n] <- valueShape v = Just (v, n)
      inspect _ = Nothing
  plottable _ = Nothing
  
  withGnuplotData ::
    [(T.Text, T.Text)] ->
    [(T.Text, [Value])] ->
    ([T.Text] -> [T.Text] -> ScriptM a) ->
    ScriptM a
  withGnuplotData sets [] cont = uncurry cont $ unzip $ reverse sets
  withGnuplotData sets ((f, vs) : xys) cont =
    withTempFile $ \fname -> do
      liftIO $ T.writeFile fname $ formatDataForGnuplot vs
      withGnuplotData ((f, f <> "='" <> T.pack fname <> "'") : sets) xys cont
  
  loadBMP :: FilePath -> ScriptM (Compound Value)
  loadBMP bmpfile = do
    res <- liftIO $ BMP.readBMP bmpfile
    case res of
      Left err ->
        throwError $ "Failed to read BMP:\n" <> showText err
      Right bmp -> do
        let bmp_bs = BMP.unpackBMPToRGBA32 bmp
            (w, h) = BMP.bmpDimensions bmp
            shape = SVec.fromList [fromIntegral h, fromIntegral w]
            pix l =
              let (i, j) = l `divMod` w
                  l' = (h - 1 - i) * w + j
                  r = fromIntegral $ bmp_bs `BS.index` (l' * 4)
                  g = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 1)
                  b = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 2)
                  a = fromIntegral $ bmp_bs `BS.index` (l' * 4 + 3)
               in (a `shiftL` 24) .|. (r `shiftL` 16) .|. (g `shiftL` 8) .|. b
        pure $ ValueAtom $ U32Value shape $ SVec.generate (w * h) pix
  
  loadImage :: FilePath -> ScriptM (Compound Value)
  loadImage imgfile =
    withTempDir $ \dir -> do
      let bmpfile = dir </> takeBaseName imgfile `replaceExtension` "bmp"
      void $ system "convert" [imgfile, "-type", "TrueColorAlpha", bmpfile] mempty
      loadBMP bmpfile
  
  loadPCM :: Int -> FilePath -> ScriptM (Compound Value)
  loadPCM num_channels pcmfile = do
    contents <- liftIO $ LBS.readFile pcmfile
    let v = SVec.byteStringToVector $ LBS.toStrict contents
        channel_length = SVec.length v `div` num_channels
        shape =
          SVec.fromList
            [ fromIntegral num_channels,
              fromIntegral channel_length
            ]
        -- ffmpeg outputs audio data in column-major format. `backPermuter` computes the
        -- tranposed indexes for a backpermutation.
        backPermuter i = (i `mod` channel_length) * num_channels + i `div` channel_length
        perm = SVec.generate (SVec.length v) backPermuter
    pure $ ValueAtom $ F64Value shape $ SVec.backpermute v perm
  
  loadAudio :: FilePath -> ScriptM (Compound Value)
  loadAudio audiofile = do
    s <- system "ffprobe" [audiofile, "-show_entries", "stream=channels", "-select_streams", "a", "-of", "compact=p=0:nk=1", "-v", "0"] mempty
    case T.decimal s of
      Right (num_channels, _) -> do
        withTempDir $ \dir -> do
          let pcmfile = dir </> takeBaseName audiofile `replaceExtension` "pcm"
          void $ system "ffmpeg" ["-i", audiofile, "-c:a", "pcm_f64le", "-map", "0", "-f", "data", pcmfile] mempty
          loadPCM num_channels pcmfile
      _ -> throwError "$loadImg failed to detect the number of channels in the audio input"
  
  literateBuiltin :: EvalBuiltin ScriptM
  literateBuiltin server "loadimg" vs
    | [v] <- vs = do
        v' <- getHaskellValue server v
        case v' of
          Just path -> do
            let path' = map (chr . fromIntegral) (path :: [Word8])
            valToExpValue <$> loadImage path'
          _ -> bad
    | otherwise = bad
    where
      bad =
        throwError $
          "$loadimg does not accept arguments of types: "
            <> T.intercalate ", " (map (prettyText . fmap scriptValueType) vs)
  literateBuiltin server "loadaudio" vs
    | [v] <- vs = do
        v' <- getHaskellValue server v
        case v' of
          Just path -> do
            let path' = map (chr . fromIntegral) (path :: [Word8])
            valToExpValue <$> loadAudio path'
          _ -> bad
    | otherwise = bad
    where
      bad =
        throwError $
          "$loadaudio does not accept arguments of types: "
            <> T.intercalate ", " (map (prettyText . fmap scriptValueType) vs)
  literateBuiltin server f vs =
    scriptBuiltin "." server f vs
  
  -- | Some of these only make sense for @futhark literate@, but enough
  -- are also sensible for @futhark script@ that we can share them.
  data Options = Options
    { scriptBackend :: String,
      scriptFuthark :: Maybe FilePath,
      scriptExtraOptions :: [String],
      scriptCompilerOptions :: [String],
      scriptSkipCompilation :: Bool,
      scriptOutput :: Maybe FilePath,
      scriptVerbose :: Int,
      scriptStopOnError :: Bool,
      scriptBinary :: Bool,
      scriptExps :: [Either FilePath T.Text]
    }
  
  -- | The configuration before any user-provided options are processed.
  initialOptions :: Options
  initialOptions =
    Options
      { scriptBackend = "c",
        scriptFuthark = Nothing,
        scriptExtraOptions = [],
        scriptCompilerOptions = [],
        scriptSkipCompilation = False,
        scriptOutput = Nothing,
        scriptVerbose = 0,
        scriptStopOnError = False,
        scriptBinary = False,
        scriptExps = []
      }
  
  data Env = Env
    { envImgDir :: FilePath,
      envOpts :: Options,
      envServer :: ScriptServer,
      envHash :: T.Text
    }
  
  newFile :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM FilePath
  newFile env (fname_desired, template) m = do
    let fname_base = fromMaybe (T.unpack (envHash env) <> "-" <> template) fname_desired
        fname = envImgDir env </> fname_base
    exists <- liftIO $ doesFileExist fname
    liftIO $ ensureCacheDirectory $ envImgDir env
    when (exists && scriptVerbose (envOpts env) > 0) $
      liftIO . T.hPutStrLn stderr $
        "Using existing file: " <> T.pack fname
    unless exists $ do
      when (scriptVerbose (envOpts env) > 0) $
        liftIO . T.hPutStrLn stderr $
          "Generating new file: " <> T.pack fname
      m fname
    modify $ \s -> s {stateFiles = S.insert fname $ stateFiles s}
    pure fname
  
  newFileContents :: Env -> (Maybe FilePath, FilePath) -> (FilePath -> ScriptM ()) -> ScriptM T.Text
  newFileContents env f m =
    liftIO . T.readFile =<< newFile env f m
  
  processDirective :: Env -> Directive -> ScriptM T.Text
  processDirective env (DirectiveBrief d) =
    processDirective env d
  processDirective env (DirectiveCovert d) =
    processDirective env d
  processDirective env (DirectiveRes e) = do
    result <-
      newFileContents env (Nothing, "eval.txt") $ \resultf -> do
        v <- either nope pure =<< evalExpToGround literateBuiltin (envServer env) e
        liftIO $ T.writeFile resultf $ prettyText v
    pure $ T.unlines ["```", result, "```"]
    where
      nope t =
        throwError $ "Cannot show value of type " <> prettyText t
  --
  processDirective env (DirectiveImg e params) = do
    fmap imgBlock . newFile env (imgFile params, "img.png") $ \pngfile -> do
      maybe_v <- evalExpToGround literateBuiltin (envServer env) e
      case maybe_v of
        Right (ValueAtom v)
          | Just bmp <- valueToBMP v -> do
              withTempDir $ \dir -> do
                let bmpfile = dir </> "img.bmp"
                liftIO $ LBS.writeFile bmpfile bmp
                void $ system "convert" [bmpfile, pngfile] mempty
        Right v ->
          nope $ fmap valueType v
        Left t ->
          nope t
    where
      nope t =
        throwError $
          "Cannot create image from value of type " <> prettyText t
  --
  processDirective env (DirectivePlot e size) = do
    fmap imgBlock . newFile env (Nothing, "plot.png") $ \pngfile -> do
      maybe_v <- evalExpToGround literateBuiltin (envServer env) e
      case maybe_v of
        Right v
          | Just vs <- plottable2d v ->
              plotWith [(Nothing, vs)] pngfile
        Right (ValueRecord m)
          | Just m' <- traverse plottable2d m -> do
              plotWith (map (first Just) $ M.toList m') pngfile
        Right v ->
          throwError $ "Cannot plot value of type " <> prettyText (fmap valueType v)
        Left t ->
          throwError $ "Cannot plot opaque value of type " <> prettyText t
    where
      plottable2d v = do
        [x, y] <- plottable v
        Just [x, y]
  
      tag (Nothing, xys) j = ("data" <> showText (j :: Int), xys)
      tag (Just f, xys) _ = (f, xys)
  
      plotWith xys pngfile =
        withGnuplotData [] (zipWith tag xys [0 ..]) $ \fs sets -> do
          let size' = T.pack $
                case size of
                  Nothing -> "500,500"
                  Just (w, h) -> show w ++ "," ++ show h
              plotCmd f title =
                let title' = case title of
                      Nothing -> "notitle"
                      Just x -> "title '" <> x <> "'"
                 in f <> " " <> title' <> " with lines"
              cmds = T.intercalate ", " (zipWith plotCmd fs (map fst xys))
              script =
                T.unlines
                  [ "set terminal png size " <> size' <> " enhanced",
                    "set output '" <> T.pack pngfile <> "'",
                    "set key outside",
                    T.unlines sets,
                    "plot " <> cmds
                  ]
          void $ system "gnuplot" [] script
  --
  processDirective env (DirectiveGnuplot e script) = do
    fmap imgBlock . newFile env (Nothing, "plot.png") $ \pngfile -> do
      maybe_v <- evalExpToGround literateBuiltin (envServer env) e
      case maybe_v of
        Right (ValueRecord m)
          | Just m' <- traverse plottable m ->
              plotWith (M.toList m') pngfile
        Right v ->
          throwError $ "Cannot plot value of type " <> prettyText (fmap valueType v)
        Left t ->
          throwError $ "Cannot plot opaque value of type " <> prettyText t
    where
      plotWith xys pngfile = withGnuplotData [] xys $ \_ sets -> do
        let script' =
              T.unlines
                [ "set terminal png enhanced",
                  "set output '" <> T.pack pngfile <> "'",
                  T.unlines sets,
                  script
                ]
        void $ system "gnuplot" [] script'
  --
  processDirective env (DirectiveVideo e params) = do
    unless (format `elem` ["webm", "gif"]) $
      throwError $
        "Unknown video format: " <> format
  
    let file = (videoFile params, "video" <.> T.unpack format)
    fmap (videoBlock params) . newFile env file $ \videofile -> do
      v <- evalExp literateBuiltin (envServer env) e
      let nope =
            throwError $
              "Cannot produce video from value of type " <> prettyText (fmap scriptValueType v)
      case v of
        ValueAtom SValue {} -> do
          ValueAtom arr <- getExpValue (envServer env) v
          case valueToBMPs arr of
            Nothing -> nope
            Just bmps ->
              withTempDir $ \dir -> do
                zipWithM_ (writeBMPFile dir) [0 ..] bmps
                onWebM videofile =<< bmpsToVideo dir
        ValueTuple [stepfun, initial, num_frames]
          | ValueAtom (SFun stepfun' _ [_, _] closure) <- stepfun,
            ValueAtom (SValue "i64" _) <- num_frames -> do
              Just (ValueAtom num_frames') <-
                mapM getValue <$> getExpValue (envServer env) num_frames
              withTempDir $ \dir -> do
                let num_frames_int = fromIntegral (num_frames' :: Int64)
                renderFrames dir (stepfun', map ValueAtom closure) initial num_frames_int
                onWebM videofile =<< bmpsToVideo dir
        _ ->
          nope
    where
      framerate = fromMaybe 30 $ videoFPS params
      format = fromMaybe "webm" $ videoFormat params
      bmpfile dir j = dir </> printf "frame%010d.bmp" (j :: Int)
  
      (progressStep, progressDone)
        | fancyTerminal,
          scriptVerbose (envOpts env) > 0 =
            ( \j num_frames -> liftIO $ do
                T.putStr $
                  "\r"
                    <> progressBar
                      (ProgressBar 40 (fromIntegral num_frames - 1) (fromIntegral j))
                    <> "generating frame "
                    <> prettyText (j + 1)
                    <> "/"
                    <> prettyText num_frames
                    <> " "
                hFlush stdout,
              liftIO $ T.putStrLn ""
            )
        | otherwise =
            (\_ _ -> pure (), pure ())
  
      renderFrames dir (stepfun, closure) initial num_frames = do
        foldM_ frame initial [0 .. num_frames - 1]
        progressDone
        where
          frame old_state j = do
            progressStep j num_frames
            v <-
              evalExp literateBuiltin (envServer env)
                . Call (FuncFut stepfun)
                . map valueToExp
                $ closure ++ [old_state]
            freeValue (envServer env) old_state
  
            let nope =
                  throwError $
                    "Cannot handle step function return type: "
                      <> prettyText (fmap scriptValueType v)
  
            case v of
              ValueTuple [arr_v@(ValueAtom SValue {}), new_state] -> do
                ValueAtom arr <- getExpValue (envServer env) arr_v
                freeValue (envServer env) arr_v
                case valueToBMP arr of
                  Nothing -> nope
                  Just bmp -> do
                    writeBMPFile dir j bmp
                    pure new_state
              _ -> nope
  
      writeBMPFile dir j bmp =
        liftIO $ LBS.writeFile (bmpfile dir j) bmp
  
      bmpsToVideo dir = do
        void $
          system
            "ffmpeg"
            [ "-y",
              "-r",
              show framerate,
              "-i",
              dir </> "frame%010d.bmp",
              "-c:v",
              "libvpx-vp9",
              "-pix_fmt",
              "yuv420p",
              "-b:v",
              "2M",
              dir </> "video.webm"
            ]
            mempty
        pure $ dir </> "video.webm"
  
      onWebM videofile webmfile
        | format == "gif" =
            void $ system "ffmpeg" ["-i", webmfile, videofile] mempty
        | otherwise =
            liftIO $ copyFile webmfile videofile
  
  --
  processDirective env (DirectiveAudio e params) = do
    fmap imgBlock . newFile env (Nothing, "output." <> T.unpack output_format) $
      \audiofile -> do
        withTempDir $ \dir -> do
          maybe_v <- evalExpToGround literateBuiltin (envServer env) e
          maybe_raw_files <- toRawFiles dir maybe_v
          case maybe_raw_files of
            (input_format, raw_files) -> do
              void $
                system
                  "ffmpeg"
                  ( concatMap
                      ( \raw_file ->
                          [ "-f",
                            input_format,
                            "-ar",
                            show sampling_frequency,
                            "-i",
                            raw_file
                          ]
                      )
                      raw_files
                      ++ [ "-f",
                           T.unpack output_format,
                           "-filter_complex",
                           concatMap
                             (\i -> "[" <> show i <> ":a]")
                             [0 .. length raw_files - 1]
                             <> "amerge=inputs="
                             <> show (length raw_files)
                             <> "[a]",
                           "-map",
                           "[a]",
                           audiofile
                         ]
                  )
                  mempty
    where
      writeRaw dir name v = do
        let rawfile = dir </> name
        let Just bytes = toBytes v
        liftIO $ LBS.writeFile rawfile $ LBS.fromStrict bytes
  
      toRawFiles dir (Right (ValueAtom v))
        | length (valueShape v) == 1,
          Just input_format <- toFfmpegFormat v = do
            writeRaw dir "raw.pcm" v
            pure (input_format, [dir </> "raw.pcm"])
        | length (valueShape v) == 2,
          Just input_format <- toFfmpegFormat v = do
            (input_format,)
              <$> zipWithM
                ( \v' i -> do
                    let file_name = "raw-" <> show i <> ".pcm"
                    writeRaw dir file_name v'
                    pure $ dir </> file_name
                )
                (valueElems v)
                [0 :: Int ..]
      toRawFiles _ v = nope $ fmap (fmap valueType) v
  
      toFfmpegFormat I8Value {} = Just "s8"
      toFfmpegFormat U8Value {} = Just "u8"
      toFfmpegFormat I16Value {} = Just "s16le"
      toFfmpegFormat U16Value {} = Just "u16le"
      toFfmpegFormat I32Value {} = Just "s32le"
      toFfmpegFormat U32Value {} = Just "u32le"
      toFfmpegFormat F32Value {} = Just "f32le"
      toFfmpegFormat F64Value {} = Just "f64le"
      toFfmpegFormat _ = Nothing
  
      toBytes (I8Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (U8Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (I16Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (U16Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (I32Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (U32Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (F32Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes (F64Value _ bytes) = Just $ SVec.vectorToByteString bytes
      toBytes _ = Nothing
  
      output_format = fromMaybe "wav" $ audioCodec params
      sampling_frequency = fromMaybe 44100 $ audioSamplingFrequency params
      nope _ = throwError "Cannot create audio from value"
  
  -- Did this script block succeed or fail?
  data Failure = Failure | Success
    deriving (Eq, Ord, Show)
  
  processBlock :: Env -> Block -> IO (Failure, T.Text, Files)
  processBlock _ (BlockCode code)
    | T.null code = pure (Success, mempty, mempty)
    | otherwise = pure (Success, "```futhark\n" <> code <> "```\n", mempty)
  processBlock _ (BlockComment pretty) =
    pure (Success, pretty, mempty)
  processBlock env (BlockDirective directive text) = do
    when (scriptVerbose (envOpts env) > 0) $
      T.hPutStrLn stderr . PP.docText $
        "Processing " <> PP.align (PP.pretty directive) <> "..."
    let prompt = case directive of
          DirectiveCovert _ -> mempty
          DirectiveBrief _ ->
            "```\n" <> PP.docText (pprDirective False directive) <> "\n```\n"
          _ ->
            "```\n" <> text <> "```\n"
        env' = env {envHash = hashText (envHash env <> prettyText directive)}
    (r, files) <- runScriptM $ processDirective env' directive
    case r of
      Left err -> failed prompt err files
      Right t -> pure (Success, prompt <> "\n" <> t, files)
    where
      failed prompt err files = do
        let message = prettyTextOneLine directive <> " failed:\n" <> err <> "\n"
        liftIO $ T.hPutStr stderr message
        when (scriptStopOnError (envOpts env)) exitFailure
        pure
          ( Failure,
            T.unlines [prompt, "**FAILED**", "```", err, "```"],
            files
          )
  
  -- Delete all files in the given directory that are not contained in
  -- 'files'.
  cleanupImgDir :: Env -> Files -> IO ()
  cleanupImgDir env keep_files =
    mapM_ toRemove . filter (not . (`S.member` keep_files))
      =<< (directoryContents (envImgDir env) `catchError` onError)
    where
      onError e
        | isDoesNotExistError e = pure []
        | otherwise = throwError e
      toRemove f = do
        when (scriptVerbose (envOpts env) > 0) $
          T.hPutStrLn stderr $
            "Deleting unused file: " <> T.pack f
        removePathForcibly f
  
  processScript :: Env -> [Block] -> IO (Failure, T.Text)
  processScript env script = do
    (failures, outputs, files) <-
      unzip3 <$> mapM (processBlock env) script
    cleanupImgDir env $
      (envImgDir env </> "CACHEDIR.TAG") `S.insert` mconcat files
    pure (L.foldl' min Success failures, T.intercalate "\n" outputs)
  
  -- | Common command line options that transform 'Options'.
  scriptCommandLineOptions :: [FunOptDescr Options]
  scriptCommandLineOptions =
    [ Option
        []
        ["backend"]
        ( ReqArg
            (\backend -> Right $ \config -> config {scriptBackend = backend})
            "PROGRAM"
        )
        "The compiler used (defaults to 'c').",
      Option
        []
        ["futhark"]
        ( ReqArg
            (\prog -> Right $ \config -> config {scriptFuthark = Just prog})
            "PROGRAM"
        )
        "The binary used for operations (defaults to same binary as 'futhark script').",
      Option
        "p"
        ["pass-option"]
        ( ReqArg
            ( \opt ->
                Right $ \config ->
                  config {scriptExtraOptions = opt : scriptExtraOptions config}
            )
            "OPT"
        )
        "Pass this option to programs being run.",
      Option
        []
        ["pass-compiler-option"]
        ( ReqArg
            ( \opt ->
                Right $ \config ->
                  config {scriptCompilerOptions = opt : scriptCompilerOptions config}
            )
            "OPT"
        )
        "Pass this option to the compiler.",
      Option
        []
        ["skip-compilation"]
        (NoArg $ Right $ \config -> config {scriptSkipCompilation = True})
        "Use already compiled program.",
      Option
        "v"
        ["verbose"]
        (NoArg $ Right $ \config -> config {scriptVerbose = scriptVerbose config + 1})
        "Enable logging. Pass multiple times for more."
    ]
  
  commandLineOptions :: [FunOptDescr Options]
  commandLineOptions =
    scriptCommandLineOptions
      <> [ Option
             "o"
             ["output"]
             (ReqArg (\opt -> Right $ \config -> config {scriptOutput = Just opt}) "FILE")
             "Override output file. Image directory is set to basename appended with -img/.",
           Option
             []
             ["stop-on-error"]
             (NoArg $ Right $ \config -> config {scriptStopOnError = True})
             "Stop and do not produce output file if any directive fails."
         ]
  
  -- | Start up (and eventually shut down) a Futhark server
  -- corresponding to the provided program. If the program has a @.fut@
  -- extension, it will be compiled automatically.
  prepareServer :: FilePath -> Options -> (ScriptServer -> IO a) -> IO a
  prepareServer prog opts f = do
    futhark <- maybe getExecutablePath pure $ scriptFuthark opts
  
    let is_fut = takeExtension prog == ".fut"
  
    unless (scriptSkipCompilation opts || not is_fut) $ do
      let compile_options = "--server" : scriptCompilerOptions opts
      when (scriptVerbose opts > 0) $
        T.hPutStrLn stderr $
          "Compiling " <> T.pack prog <> "..."
      when (scriptVerbose opts > 1) $
        T.hPutStrLn stderr $
          T.pack $
            unwords compile_options
  
      let onError err = do
            T.hPutStrLn stderr err
            exitFailure
  
      void $
        either onError pure <=< runExceptT $
          compileProgram compile_options (FutharkExe futhark) (scriptBackend opts) prog
  
    let run_options = scriptExtraOptions opts
        onLine "call" l = T.putStrLn l
        onLine "startup" l = T.putStrLn l
        onLine _ _ = pure ()
        prog' = if is_fut then dropExtension prog else prog
        cfg =
          (futharkServerCfg ("." </> prog') run_options)
            { cfgOnLine =
                if scriptVerbose opts > 0
                  then onLine
                  else const . const $ pure ()
            }
  
    withScriptServer cfg f
  
  -- | Run @futhark literate@.
  main :: String -> [String] -> IO ()
  main = mainWithOptions initialOptions commandLineOptions "program" $ \args opts ->
    case args of
      [prog] -> Just $ do
        futhark <- maybe getExecutablePath pure $ scriptFuthark opts
        let onError err = do
              T.hPutStrLn stderr err
              exitFailure
        proghash <-
          either onError pure <=< runExceptT $
            system futhark ["hash", prog] mempty
        script <- parseProgFile prog
  
        orig_dir <- getCurrentDirectory
        let entryOpt v = "--entry-point=" ++ T.unpack v
            opts' =
              opts
                { scriptCompilerOptions =
                    map entryOpt (S.toList (varsInScripts script))
                      <> scriptCompilerOptions opts
                }
        prepareServer prog opts' $ \server -> do
          let mdfile = fromMaybe (prog `replaceExtension` "md") $ scriptOutput opts
              prog_dir = takeDirectory prog
              imgdir = dropExtension (takeFileName mdfile) <> "-img"
              env =
                Env
                  { envServer = server,
                    envOpts = opts,
                    envHash = proghash,
                    envImgDir = imgdir
                  }
  
          when (scriptVerbose opts > 0) $ do
            T.hPutStrLn stderr $ "Executing from " <> T.pack prog_dir
          setCurrentDirectory prog_dir
  
          (failure, md) <- processScript env script
          T.writeFile (orig_dir </> mdfile) md
          when (failure == Failure) exitFailure
      _ -> Nothing
  