module Nbparts.Unpack.Outputs where

import Control.Monad.Error.Class (MonadError, throwError)
import Control.Monad.State.Strict (MonadState)
import Control.Monad.State.Strict qualified as State
import Data.ByteString (ByteString)
import Data.Ipynb qualified as Ipynb
import Data.List qualified as List
import Data.Map qualified as Map
import Data.Maybe qualified as Maybe
import Data.Text (Text)
import Nbparts.Types
  ( UnembeddedCellOutput (DisplayData, Err, ExecuteResult, Stream),
    UnembeddedNotebookOutputs (UnembeddedNotebookOutputs),
    UnpackError (UnpackMissingCellIdError),
  )
import Nbparts.Unpack.Mime (unembedMimeBundle)

collectOutputs ::
  FilePath ->
  Ipynb.Notebook a ->
  Either UnpackError (UnembeddedNotebookOutputs, [(FilePath, ByteString)])
collectOutputs :: forall a.
FilePath
-> Notebook a
-> Either
     UnpackError (UnembeddedNotebookOutputs, [(FilePath, ByteString)])
collectOutputs FilePath
subdir (Ipynb.Notebook JSONMeta
_meta (Int, Int)
_format [Cell a]
cells) = StateT
  [(FilePath, ByteString)]
  (Either UnpackError)
  UnembeddedNotebookOutputs
-> [(FilePath, ByteString)]
-> Either
     UnpackError (UnembeddedNotebookOutputs, [(FilePath, ByteString)])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT StateT
  [(FilePath, ByteString)]
  (Either UnpackError)
  UnembeddedNotebookOutputs
forall (m :: * -> *).
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
m UnembeddedNotebookOutputs
nbOutputsState []
  where
    nbOutputsState ::
      ( MonadState [(FilePath, ByteString)] m,
        MonadError UnpackError m
      ) =>
      m UnembeddedNotebookOutputs
    nbOutputsState :: forall (m :: * -> *).
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
m UnembeddedNotebookOutputs
nbOutputsState = do
      let codeOutputs :: [(Maybe Text, [Output a])]
codeOutputs = (Cell a -> Maybe (Maybe Text, [Output a]))
-> [Cell a] -> [(Maybe Text, [Output a])]
forall a b. (a -> Maybe b) -> [a] -> [b]
Maybe.mapMaybe Cell a -> Maybe (Maybe Text, [Output a])
forall a. Cell a -> Maybe (Maybe Text, [Output a])
extractCodeOutputs [Cell a]
cells
      [(Text, [UnembeddedCellOutput])]
entries <- ((Maybe Text, [Output a]) -> m (Text, [UnembeddedCellOutput]))
-> [(Maybe Text, [Output a])] -> m [(Text, [UnembeddedCellOutput])]
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 ((Maybe Text -> [Output a] -> m (Text, [UnembeddedCellOutput]))
-> (Maybe Text, [Output a]) -> m (Text, [UnembeddedCellOutput])
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry (FilePath
-> Maybe Text -> [Output a] -> m (Text, [UnembeddedCellOutput])
forall (m :: * -> *) a.
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
FilePath
-> Maybe Text -> [Output a] -> m (Text, [UnembeddedCellOutput])
unembedCodeOutputs FilePath
subdir)) [(Maybe Text, [Output a])]
codeOutputs
      UnembeddedNotebookOutputs -> m UnembeddedNotebookOutputs
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Map Text [UnembeddedCellOutput] -> UnembeddedNotebookOutputs
UnembeddedNotebookOutputs (Map Text [UnembeddedCellOutput] -> UnembeddedNotebookOutputs)
-> Map Text [UnembeddedCellOutput] -> UnembeddedNotebookOutputs
forall a b. (a -> b) -> a -> b
$ [(Text, [UnembeddedCellOutput])] -> Map Text [UnembeddedCellOutput]
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, [UnembeddedCellOutput])]
entries)

    extractCodeOutputs :: Ipynb.Cell a -> Maybe (Maybe Text, [Ipynb.Output a])
    extractCodeOutputs :: forall a. Cell a -> Maybe (Maybe Text, [Output a])
extractCodeOutputs (Ipynb.Cell (Ipynb.Code Maybe Int
_ [Output a]
outputs) Maybe Text
maybeCellId Source
_ JSONMeta
_ Maybe MimeAttachments
_)
      | [Output a] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
List.null [Output a]
outputs = Maybe (Maybe Text, [Output a])
forall a. Maybe a
Nothing -- Filter out cells with empty outputs to avoid serialising empty lists.
      | Bool
otherwise = (Maybe Text, [Output a]) -> Maybe (Maybe Text, [Output a])
forall a. a -> Maybe a
Just (Maybe Text
maybeCellId, [Output a]
outputs)
    extractCodeOutputs Cell a
_ = Maybe (Maybe Text, [Output a])
forall a. Maybe a
Nothing

unembedCodeOutputs ::
  (MonadState [(FilePath, ByteString)] m, MonadError UnpackError m) =>
  FilePath ->
  Maybe Text ->
  [Ipynb.Output a] ->
  m (Text, [UnembeddedCellOutput])
unembedCodeOutputs :: forall (m :: * -> *) a.
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
FilePath
-> Maybe Text -> [Output a] -> m (Text, [UnembeddedCellOutput])
unembedCodeOutputs FilePath
subdir (Just Text
cellId) [Output a]
outputs = (Text
cellId,) ([UnembeddedCellOutput] -> (Text, [UnembeddedCellOutput]))
-> m [UnembeddedCellOutput] -> m (Text, [UnembeddedCellOutput])
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Output a -> m UnembeddedCellOutput)
-> [Output a] -> m [UnembeddedCellOutput]
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 (FilePath -> Output a -> m UnembeddedCellOutput
forall (m :: * -> *) a.
MonadState [(FilePath, ByteString)] m =>
FilePath -> Output a -> m UnembeddedCellOutput
unembedOutput FilePath
subdir) [Output a]
outputs
unembedCodeOutputs FilePath
_ Maybe Text
Nothing [Output a]
_ = UnpackError -> m (Text, [UnembeddedCellOutput])
forall a. UnpackError -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError UnpackError
UnpackMissingCellIdError

unembedOutput :: (MonadState [(FilePath, ByteString)] m) => FilePath -> Ipynb.Output a -> m UnembeddedCellOutput
unembedOutput :: forall (m :: * -> *) a.
MonadState [(FilePath, ByteString)] m =>
FilePath -> Output a -> m UnembeddedCellOutput
unembedOutput FilePath
_ (Ipynb.Stream Text
streamName (Ipynb.Source [Text]
streamText)) = UnembeddedCellOutput -> m UnembeddedCellOutput
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UnembeddedCellOutput -> m UnembeddedCellOutput)
-> UnembeddedCellOutput -> m UnembeddedCellOutput
forall a b. (a -> b) -> a -> b
$ Text -> [Text] -> UnembeddedCellOutput
Stream Text
streamName [Text]
streamText
unembedOutput FilePath
subdir (Ipynb.DisplayData MimeBundle
displayData JSONMeta
metadata) = do
  UnembeddedMimeBundle
uDisplayData <- FilePath -> MimeBundle -> m UnembeddedMimeBundle
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundle FilePath
subdir MimeBundle
displayData
  UnembeddedCellOutput -> m UnembeddedCellOutput
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UnembeddedCellOutput -> m UnembeddedCellOutput)
-> UnembeddedCellOutput -> m UnembeddedCellOutput
forall a b. (a -> b) -> a -> b
$ UnembeddedMimeBundle -> JSONMeta -> UnembeddedCellOutput
DisplayData UnembeddedMimeBundle
uDisplayData JSONMeta
metadata
unembedOutput FilePath
subdir (Ipynb.ExecuteResult Int
executeCount MimeBundle
executeData JSONMeta
metadata) = do
  UnembeddedMimeBundle
uExData <- FilePath -> MimeBundle -> m UnembeddedMimeBundle
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundle FilePath
subdir MimeBundle
executeData
  UnembeddedCellOutput -> m UnembeddedCellOutput
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UnembeddedCellOutput -> m UnembeddedCellOutput)
-> UnembeddedCellOutput -> m UnembeddedCellOutput
forall a b. (a -> b) -> a -> b
$ Int -> UnembeddedMimeBundle -> JSONMeta -> UnembeddedCellOutput
ExecuteResult Int
executeCount UnembeddedMimeBundle
uExData JSONMeta
metadata
unembedOutput FilePath
_ (Ipynb.Err Text
errName Text
errValue [Text]
errTraceback) = UnembeddedCellOutput -> m UnembeddedCellOutput
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (UnembeddedCellOutput -> m UnembeddedCellOutput)
-> UnembeddedCellOutput -> m UnembeddedCellOutput
forall a b. (a -> b) -> a -> b
$ Text -> Text -> [Text] -> UnembeddedCellOutput
Err Text
errName Text
errValue [Text]
errTraceback