module Nbparts.Pack.Outputs where

import Control.Monad.Error.Class (MonadError, liftEither)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Coerce (coerce)
import Data.Ipynb qualified as Ipynb
import Data.Map qualified as Map
import Data.Maybe qualified as Maybe
import Nbparts.Pack.Mime (embedMimeBundle)
import Nbparts.Types
  ( NotebookOutputs (NotebookOutputs),
    PackError (PackMissingCellIdError),
    UnembeddedCellOutput (DisplayData, Err, ExecuteResult, Stream),
    UnembeddedNotebookOutputs (UnembeddedNotebookOutputs),
  )

fillOutputs ::
  (MonadError PackError m, MonadIO m) =>
  FilePath ->
  UnembeddedNotebookOutputs ->
  Ipynb.Notebook a ->
  m (Ipynb.Notebook a)
fillOutputs :: forall (m :: * -> *) a.
(MonadError PackError m, MonadIO m) =>
FilePath
-> UnembeddedNotebookOutputs -> Notebook a -> m (Notebook a)
fillOutputs FilePath
readDir UnembeddedNotebookOutputs
unembeddedOutputs (Ipynb.Notebook JSONMeta
meta (Int, Int)
format [Cell a]
cells) = do
  NotebookOutputs a
outputs <- IO (NotebookOutputs a) -> m (NotebookOutputs a)
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (NotebookOutputs a) -> m (NotebookOutputs a))
-> IO (NotebookOutputs a) -> m (NotebookOutputs a)
forall a b. (a -> b) -> a -> b
$ FilePath -> UnembeddedNotebookOutputs -> IO (NotebookOutputs a)
forall a.
FilePath -> UnembeddedNotebookOutputs -> IO (NotebookOutputs a)
embedOutputs FilePath
readDir UnembeddedNotebookOutputs
unembeddedOutputs
  [Cell a]
filledCells <- Either PackError [Cell a] -> m [Cell a]
forall e (m :: * -> *) a. MonadError e m => Either e a -> m a
liftEither (Either PackError [Cell a] -> m [Cell a])
-> Either PackError [Cell a] -> m [Cell a]
forall a b. (a -> b) -> a -> b
$ (Cell a -> Either PackError (Cell a))
-> [Cell a] -> Either PackError [Cell a]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse (NotebookOutputs a -> Cell a -> Either PackError (Cell a)
forall a. NotebookOutputs a -> Cell a -> Either PackError (Cell a)
fillCellOutputs NotebookOutputs a
outputs) [Cell a]
cells
  Notebook a -> m (Notebook a)
forall a. a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return (Notebook a -> m (Notebook a)) -> Notebook a -> m (Notebook a)
forall a b. (a -> b) -> a -> b
$ JSONMeta -> (Int, Int) -> [Cell a] -> Notebook a
forall a. JSONMeta -> (Int, Int) -> [Cell a] -> Notebook a
Ipynb.Notebook JSONMeta
meta (Int, Int)
format [Cell a]
filledCells

fillCellOutputs :: NotebookOutputs a -> Ipynb.Cell a -> Either PackError (Ipynb.Cell a)
fillCellOutputs :: forall a. NotebookOutputs a -> Cell a -> Either PackError (Cell a)
fillCellOutputs (NotebookOutputs Map Text [Output a]
outputs) (Ipynb.Cell (Ipynb.Code Maybe Int
codeExecutionCount [Output a]
_) Maybe Text
maybeCellId Source
source JSONMeta
meta Maybe MimeAttachments
attachments) = do
  Text
cellId <- case Maybe Text
maybeCellId of
    Just Text
cellId -> Text -> Either PackError Text
forall a b. b -> Either a b
Right Text
cellId
    Maybe Text
Nothing -> PackError -> Either PackError Text
forall a b. a -> Either a b
Left PackError
PackMissingCellIdError

  -- If we can't find the outputs, assume that there are no outputs.
  let cellOutputs :: [Output a]
cellOutputs = [Output a] -> Maybe [Output a] -> [Output a]
forall a. a -> Maybe a -> a
Maybe.fromMaybe [] (Maybe [Output a] -> [Output a]) -> Maybe [Output a] -> [Output a]
forall a b. (a -> b) -> a -> b
$ Text -> Map Text [Output a] -> Maybe [Output a]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
cellId Map Text [Output a]
outputs

  Cell a -> Either PackError (Cell a)
forall a. a -> Either PackError a
forall (m :: * -> *) a. Monad m => a -> m a
return (Cell a -> Either PackError (Cell a))
-> Cell a -> Either PackError (Cell a)
forall a b. (a -> b) -> a -> b
$ CellType a
-> Maybe Text
-> Source
-> JSONMeta
-> Maybe MimeAttachments
-> Cell a
forall a.
CellType a
-> Maybe Text
-> Source
-> JSONMeta
-> Maybe MimeAttachments
-> Cell a
Ipynb.Cell (Maybe Int -> [Output a] -> CellType a
forall a. Maybe Int -> [Output a] -> CellType a
Ipynb.Code Maybe Int
codeExecutionCount [Output a]
cellOutputs) (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
cellId) Source
source JSONMeta
meta Maybe MimeAttachments
attachments
fillCellOutputs NotebookOutputs a
_ Cell a
cell = Cell a -> Either PackError (Cell a)
forall a. a -> Either PackError a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Cell a
cell

embedOutputs :: FilePath -> UnembeddedNotebookOutputs -> IO (NotebookOutputs a)
embedOutputs :: forall a.
FilePath -> UnembeddedNotebookOutputs -> IO (NotebookOutputs a)
embedOutputs FilePath
readDir = (Map Text [UnembeddedCellOutput] -> IO (NotebookOutputs Any))
-> UnembeddedNotebookOutputs -> IO (NotebookOutputs a)
forall a b. Coercible a b => a -> b
coerce ((Map Text [UnembeddedCellOutput] -> IO (NotebookOutputs Any))
 -> UnembeddedNotebookOutputs -> IO (NotebookOutputs a))
-> (Map Text [UnembeddedCellOutput] -> IO (NotebookOutputs Any))
-> UnembeddedNotebookOutputs
-> IO (NotebookOutputs a)
forall a b. (a -> b) -> a -> b
$ (Map Text [Output Any] -> NotebookOutputs Any)
-> IO (Map Text [Output Any]) -> IO (NotebookOutputs Any)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Map Text [Output Any] -> NotebookOutputs Any
forall a. Map Text [Output a] -> NotebookOutputs a
NotebookOutputs (IO (Map Text [Output Any]) -> IO (NotebookOutputs Any))
-> (Map Text [UnembeddedCellOutput] -> IO (Map Text [Output Any]))
-> Map Text [UnembeddedCellOutput]
-> IO (NotebookOutputs Any)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ([UnembeddedCellOutput] -> IO [Output Any])
-> Map Text [UnembeddedCellOutput] -> IO (Map Text [Output Any])
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Map Text a -> f (Map Text b)
traverse ((UnembeddedCellOutput -> IO (Output Any))
-> [UnembeddedCellOutput] -> IO [Output Any]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
(a -> m b) -> t a -> m (t b)
forall (m :: * -> *) a b. Monad m => (a -> m b) -> [a] -> m [b]
mapM ((UnembeddedCellOutput -> IO (Output Any))
 -> [UnembeddedCellOutput] -> IO [Output Any])
-> (UnembeddedCellOutput -> IO (Output Any))
-> [UnembeddedCellOutput]
-> IO [Output Any]
forall a b. (a -> b) -> a -> b
$ FilePath -> UnembeddedCellOutput -> IO (Output Any)
forall a. FilePath -> UnembeddedCellOutput -> IO (Output a)
embedOutput FilePath
readDir)

embedOutput :: FilePath -> UnembeddedCellOutput -> IO (Ipynb.Output a)
embedOutput :: forall a. FilePath -> UnembeddedCellOutput -> IO (Output a)
embedOutput FilePath
_ (Stream Text
streamName [Text]
streamText) = Output a -> IO (Output a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Output a -> IO (Output a)) -> Output a -> IO (Output a)
forall a b. (a -> b) -> a -> b
$ Ipynb.Stream {Text
streamName :: Text
streamName :: Text
streamName, streamText :: Source
streamText = [Text] -> Source
Ipynb.Source [Text]
streamText}
embedOutput FilePath
readDir (DisplayData UnembeddedMimeBundle
unembeddedDisplayData JSONMeta
displayMetadata) = do
  MimeBundle
displayData <- FilePath -> UnembeddedMimeBundle -> IO MimeBundle
embedMimeBundle FilePath
readDir UnembeddedMimeBundle
unembeddedDisplayData
  Output a -> IO (Output a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
    Ipynb.DisplayData
      { displayData :: MimeBundle
displayData = MimeBundle
displayData,
        JSONMeta
displayMetadata :: JSONMeta
displayMetadata :: JSONMeta
displayMetadata
      }
embedOutput FilePath
readDir (ExecuteResult Int
executeCount UnembeddedMimeBundle
unembeddedExecuteData JSONMeta
executeMetadata) = do
  MimeBundle
executeData <- FilePath -> UnembeddedMimeBundle -> IO MimeBundle
embedMimeBundle FilePath
readDir UnembeddedMimeBundle
unembeddedExecuteData
  Output a -> IO (Output a)
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return
    Ipynb.ExecuteResult
      { Int
executeCount :: Int
executeCount :: Int
executeCount,
        executeData :: MimeBundle
executeData = MimeBundle
executeData,
        JSONMeta
executeMetadata :: JSONMeta
executeMetadata :: JSONMeta
executeMetadata
      }
embedOutput FilePath
_ (Err Text
errName Text
errValue [Text]
errTraceback) = Output a -> IO (Output a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Output a -> IO (Output a)) -> Output a -> IO (Output a)
forall a b. (a -> b) -> a -> b
$ Ipynb.Err {Text
errName :: Text
errName :: Text
errName, Text
errValue :: Text
errValue :: Text
errValue, [Text]
errTraceback :: [Text]
errTraceback :: [Text]
errTraceback}