module Nbparts.Unpack.Metadata where

import Data.Ipynb qualified as Ipynb
import Data.Map qualified as Map
import Data.Maybe qualified as Maybe
import Data.Text (Text)
import Nbparts.Types
  ( CellMetadata (CodeCellMetadata, GenericCellMetadata),
    NotebookMetadata (NotebookMetadata, cells, formatMajor, formatMinor, toplevel),
    UnpackError (UnpackMissingCellIdError),
  )

extractNotebookVersion :: Ipynb.Notebook a -> (Int, Int)
extractNotebookVersion :: forall a. Notebook a -> (Int, Int)
extractNotebookVersion (Ipynb.Notebook JSONMeta
_ (Int, Int)
format [Cell a]
_) = (Int, Int)
format

collectMetadata :: Ipynb.Notebook a -> Either UnpackError NotebookMetadata
collectMetadata :: forall a. Notebook a -> Either UnpackError NotebookMetadata
collectMetadata (Ipynb.Notebook JSONMeta
toplevelMeta (Int
formatMajor, Int
formatMinor) [Cell a]
cells) = do
  [(Text, CellMetadata)]
cellsMetaList <- (Cell a -> Either UnpackError (Text, CellMetadata))
-> [Cell a] -> Either UnpackError [(Text, CellMetadata)]
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 Cell a -> Either UnpackError (Text, CellMetadata)
forall a. Cell a -> Either UnpackError (Text, CellMetadata)
extractCellMetadata [Cell a]
cells
  -- Remove empty metadata for compactness.
  let filtered :: [(Text, CellMetadata)]
filtered = ((Text, CellMetadata) -> Bool)
-> [(Text, CellMetadata)] -> [(Text, CellMetadata)]
forall a. (a -> Bool) -> [a] -> [a]
filter (Bool -> Bool
not (Bool -> Bool)
-> ((Text, CellMetadata) -> Bool) -> (Text, CellMetadata) -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. CellMetadata -> Bool
isEmptyMeta (CellMetadata -> Bool)
-> ((Text, CellMetadata) -> CellMetadata)
-> (Text, CellMetadata)
-> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text, CellMetadata) -> CellMetadata
forall a b. (a, b) -> b
snd) [(Text, CellMetadata)]
cellsMetaList
  let cellsMeta :: Map Text CellMetadata
cellsMeta = [(Text, CellMetadata)] -> Map Text CellMetadata
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList [(Text, CellMetadata)]
filtered
  NotebookMetadata -> Either UnpackError NotebookMetadata
forall a. a -> Either UnpackError a
forall (m :: * -> *) a. Monad m => a -> m a
return NotebookMetadata {Int
formatMajor :: Int
formatMajor :: Int
formatMajor, Int
formatMinor :: Int
formatMinor :: Int
formatMinor, toplevel :: JSONMeta
toplevel = JSONMeta
toplevelMeta, cells :: Map Text CellMetadata
cells = Map Text CellMetadata
cellsMeta}

isEmptyMeta :: CellMetadata -> Bool
isEmptyMeta :: CellMetadata -> Bool
isEmptyMeta (CodeCellMetadata Maybe Int
exeCount (Ipynb.JSONMeta Map Text Value
meta)) = Maybe Int -> Bool
forall a. Maybe a -> Bool
Maybe.isNothing Maybe Int
exeCount Bool -> Bool -> Bool
&& Map Text Value -> Bool
forall k a. Map k a -> Bool
Map.null Map Text Value
meta
isEmptyMeta (GenericCellMetadata (Ipynb.JSONMeta Map Text Value
meta)) = Map Text Value -> Bool
forall k a. Map k a -> Bool
Map.null Map Text Value
meta

extractCellMetadata :: Ipynb.Cell a -> Either UnpackError (Text, CellMetadata)
extractCellMetadata :: forall a. Cell a -> Either UnpackError (Text, CellMetadata)
extractCellMetadata (Ipynb.Cell (Ipynb.Code Maybe Int
exeCount [Output a]
_) Maybe Text
maybeCellId Source
_ JSONMeta
meta Maybe MimeAttachments
_) = case Maybe Text
maybeCellId of
  Just Text
cellId -> (Text, CellMetadata) -> Either UnpackError (Text, CellMetadata)
forall a b. b -> Either a b
Right (Text
cellId, Maybe Int -> JSONMeta -> CellMetadata
CodeCellMetadata Maybe Int
exeCount JSONMeta
meta)
  Maybe Text
Nothing -> UnpackError -> Either UnpackError (Text, CellMetadata)
forall a b. a -> Either a b
Left UnpackError
UnpackMissingCellIdError
extractCellMetadata (Ipynb.Cell CellType a
_ Maybe Text
maybeCellId Source
_ JSONMeta
meta Maybe MimeAttachments
_) = case Maybe Text
maybeCellId of
  Just Text
cellId -> (Text, CellMetadata) -> Either UnpackError (Text, CellMetadata)
forall a b. b -> Either a b
Right (Text
cellId, JSONMeta -> CellMetadata
GenericCellMetadata JSONMeta
meta)
  Maybe Text
Nothing -> UnpackError -> Either UnpackError (Text, CellMetadata)
forall a b. a -> Either a b
Left UnpackError
UnpackMissingCellIdError