module Nbparts.Unpack.Sources where

import Control.Monad.Error.Class (MonadError, throwError)
import Control.Monad.State.Class (MonadState)
import Control.Monad.State.Strict qualified as State
import Data.ByteString (ByteString)
import Data.Ipynb qualified as Ipynb
import Nbparts.Types
  ( CellSource (CellSource),
    CellType (Code, Markdown, Raw),
    UnpackError (UnpackMissingCellIdError),
  )
import Nbparts.Unpack.Mime (unembedMimeAttachments)

collectSources :: FilePath -> Ipynb.Notebook a -> Either UnpackError ([CellSource], [(FilePath, ByteString)])
collectSources :: forall a.
FilePath
-> Notebook a
-> Either UnpackError ([CellSource], [(FilePath, ByteString)])
collectSources FilePath
subdir (Ipynb.Notebook JSONMeta
_ (Int, Int)
_ [Cell a]
cells) = StateT [(FilePath, ByteString)] (Either UnpackError) [CellSource]
-> [(FilePath, ByteString)]
-> Either UnpackError ([CellSource], [(FilePath, ByteString)])
forall s (m :: * -> *) a. StateT s m a -> s -> m (a, s)
State.runStateT ((Cell a
 -> StateT [(FilePath, ByteString)] (Either UnpackError) CellSource)
-> [Cell a]
-> StateT
     [(FilePath, ByteString)] (Either UnpackError) [CellSource]
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
-> Cell a
-> StateT [(FilePath, ByteString)] (Either UnpackError) CellSource
forall (m :: * -> *) a.
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
FilePath -> Cell a -> m CellSource
convertCell FilePath
subdir) [Cell a]
cells) []

convertCell ::
  (MonadState [(FilePath, ByteString)] m, MonadError UnpackError m) =>
  FilePath ->
  Ipynb.Cell a ->
  m CellSource
convertCell :: forall (m :: * -> *) a.
(MonadState [(FilePath, ByteString)] m,
 MonadError UnpackError m) =>
FilePath -> Cell a -> m CellSource
convertCell FilePath
subdir (Ipynb.Cell CellType a
cellType Maybe Text
maybeCellId (Ipynb.Source [Text]
source) JSONMeta
_ Maybe MimeAttachments
maybeAtts) = do
  Text
cellId <- case Maybe Text
maybeCellId of
    Just Text
cId -> Text -> m Text
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Text
cId
    Maybe Text
Nothing -> UnpackError -> m Text
forall a. UnpackError -> m a
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError UnpackError
UnpackMissingCellIdError
  Maybe UnembeddedMimeAttachments
maybeUAtts <- (MimeAttachments -> m UnembeddedMimeAttachments)
-> Maybe MimeAttachments -> m (Maybe UnembeddedMimeAttachments)
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) -> Maybe a -> f (Maybe b)
traverse (FilePath -> MimeAttachments -> m UnembeddedMimeAttachments
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
FilePath -> MimeAttachments -> m UnembeddedMimeAttachments
unembedMimeAttachments FilePath
subdir) Maybe MimeAttachments
maybeAtts
  CellSource -> m CellSource
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (CellSource -> m CellSource) -> CellSource -> m CellSource
forall a b. (a -> b) -> a -> b
$ Text
-> CellType
-> [Text]
-> Maybe UnembeddedMimeAttachments
-> CellSource
CellSource Text
cellId (CellType a -> CellType
forall a. CellType a -> CellType
convertCellType CellType a
cellType) [Text]
source Maybe UnembeddedMimeAttachments
maybeUAtts

convertCellType :: Ipynb.CellType a -> CellType
convertCellType :: forall a. CellType a -> CellType
convertCellType CellType a
Ipynb.Markdown = CellType
Markdown
convertCellType CellType a
Ipynb.Raw = CellType
Raw
convertCellType (Ipynb.Code Maybe Int
_ [Output a]
_) = CellType
Code
convertCellType (Ipynb.Heading Int
_) = FilePath -> CellType
forall a. HasCallStack => FilePath -> a
error FilePath
"Unexpected heading cell" -- Should not be possible since we only deal with V4 notebooks.