module Nbparts.Pack.Metadata where

import Data.Ipynb qualified as Ipynb
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Text (Text)
import Nbparts.Types
  ( CellMetadata (CodeCellMetadata, GenericCellMetadata),
    CellMetadataTag (CodeCellMetadataTag, GenericCellMetadataTag),
    NotebookMetadata (NotebookMetadata),
    PackError (PackCellMetadataTypeMismatch, PackMissingCellIdError, actual, expected),
    emptyCodeMetadata,
    emptyGenericMetadata,
  )

fillMetadata :: NotebookMetadata -> Ipynb.Notebook a -> Either PackError (Ipynb.Notebook a)
fillMetadata :: forall a.
NotebookMetadata -> Notebook a -> Either PackError (Notebook a)
fillMetadata (NotebookMetadata Int
formatMajor Int
formatMinor JSONMeta
nbMeta Map Text CellMetadata
cellsMeta) (Ipynb.Notebook JSONMeta
_ (Int, Int)
_ [Cell a]
cells) =
  JSONMeta -> (Int, Int) -> [Cell a] -> Notebook a
forall a. JSONMeta -> (Int, Int) -> [Cell a] -> Notebook a
Ipynb.Notebook
    JSONMeta
nbMeta
    (Int
formatMajor, Int
formatMinor)
    ([Cell a] -> Notebook a)
-> Either PackError [Cell a] -> Either PackError (Notebook a)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Either PackError [Cell a]
filledCells
  where
    filledCells :: Either PackError [Cell a]
filledCells = (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 (Map Text CellMetadata -> Cell a -> Either PackError (Cell a)
forall a.
Map Text CellMetadata -> Cell a -> Either PackError (Cell a)
fillCellMetadata Map Text CellMetadata
cellsMeta) [Cell a]
cells

fillCellMetadata :: Map Text CellMetadata -> Ipynb.Cell a -> Either PackError (Ipynb.Cell a)
fillCellMetadata :: forall a.
Map Text CellMetadata -> Cell a -> Either PackError (Cell a)
fillCellMetadata Map Text CellMetadata
metaMap (Ipynb.Cell CellType a
cellType Maybe Text
maybeCellId Source
source JSONMeta
_ Maybe MimeAttachments
attachments) = do
  Text
cellId <- case Maybe Text
maybeCellId of
    Just Text
cId -> Text -> Either PackError Text
forall a b. b -> Either a b
Right Text
cId
    Maybe Text
Nothing -> PackError -> Either PackError Text
forall a b. a -> Either a b
Left PackError
PackMissingCellIdError
  let cellMeta :: CellMetadata
cellMeta = case (Text -> Map Text CellMetadata -> Maybe CellMetadata
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup Text
cellId Map Text CellMetadata
metaMap, CellType a
cellType) of
        (Just CellMetadata
meta, CellType a
_) -> CellMetadata
meta
        -- If we can't find the metadata, assume there is none.
        (Maybe CellMetadata
Nothing, Ipynb.Code Maybe Int
_ [Output a]
_) -> CellMetadata
emptyCodeMetadata
        (Maybe CellMetadata
Nothing, CellType a
_) -> CellMetadata
emptyGenericMetadata
  case (CellType a
cellType, CellMetadata
cellMeta) of
    -- Code cell expects CodeCellMetadata
    (Ipynb.Code Maybe Int
_ [Output a]
outputs, CodeCellMetadata Maybe Int
exeCount JSONMeta
meta) ->
      Cell a -> Either PackError (Cell a)
forall a b. b -> Either a b
Right (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
exeCount [Output a]
outputs) (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
cellId) Source
source JSONMeta
meta Maybe MimeAttachments
attachments
    (Ipynb.Code Maybe Int
_ [Output a]
_, GenericCellMetadata JSONMeta
_) ->
      PackError -> Either PackError (Cell a)
forall a b. a -> Either a b
Left (PackError -> Either PackError (Cell a))
-> PackError -> Either PackError (Cell a)
forall a b. (a -> b) -> a -> b
$
        PackCellMetadataTypeMismatch
          { expected :: CellMetadataTag
expected = CellMetadataTag
CodeCellMetadataTag,
            actual :: CellMetadataTag
actual = CellMetadataTag
GenericCellMetadataTag
          }
    -- Other cell types expect GenericCellMetadata
    (CellType a
_, GenericCellMetadata JSONMeta
meta) ->
      Cell a -> Either PackError (Cell a)
forall a b. b -> Either a b
Right (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 CellType a
cellType (Text -> Maybe Text
forall a. a -> Maybe a
Just Text
cellId) Source
source JSONMeta
meta Maybe MimeAttachments
attachments
    (CellType a
_, CodeCellMetadata Maybe Int
_ JSONMeta
_) ->
      PackError -> Either PackError (Cell a)
forall a b. a -> Either a b
Left (PackError -> Either PackError (Cell a))
-> PackError -> Either PackError (Cell a)
forall a b. (a -> b) -> a -> b
$
        PackCellMetadataTypeMismatch
          { expected :: CellMetadataTag
expected = CellMetadataTag
GenericCellMetadataTag,
            actual :: CellMetadataTag
actual = CellMetadataTag
CodeCellMetadataTag
          }