module Nbparts.Unpack.Mime where

import Control.Monad.State.Strict (MonadState)
import Control.Monad.State.Strict qualified as State
import Crypto.Hash qualified as Hash
import Data.ByteArray qualified as ByteArray
import Data.ByteString (ByteString)
import Data.ByteString.Base64.URL qualified as Base64.URL
import Data.ByteString.Char8 qualified as ByteString
import Data.Coerce (coerce)
import Data.Ipynb qualified as Ipynb
import Data.Map qualified as Map
import Data.Maybe qualified as Maybe
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Nbparts.Types
  ( UnembeddedMimeAttachments (UnembeddedMimeAttachments),
    UnembeddedMimeBundle (UnembeddedMimeBundle),
    UnembeddedMimeData (BinaryData, JsonData, TextualData),
  )
import Network.Mime qualified as Mime
import System.FilePath ((</>))

unembedMimeAttachments :: (MonadState [(FilePath, ByteString)] m) => FilePath -> Ipynb.MimeAttachments -> m UnembeddedMimeAttachments
unembedMimeAttachments :: forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
FilePath -> MimeAttachments -> m UnembeddedMimeAttachments
unembedMimeAttachments FilePath
subdir MimeAttachments
mimeAtts = Map Text UnembeddedMimeBundle -> UnembeddedMimeAttachments
UnembeddedMimeAttachments (Map Text UnembeddedMimeBundle -> UnembeddedMimeAttachments)
-> m (Map Text UnembeddedMimeBundle) -> m UnembeddedMimeAttachments
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> MimeBundle -> m UnembeddedMimeBundle)
-> Map Text MimeBundle -> m (Map Text UnembeddedMimeBundle)
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey Text -> MimeBundle -> m UnembeddedMimeBundle
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
Text -> MimeBundle -> m UnembeddedMimeBundle
go (MimeAttachments -> Map Text MimeBundle
forall a b. Coercible a b => a -> b
coerce MimeAttachments
mimeAtts)
  where
    go :: (MonadState [(FilePath, ByteString)] m) => Text -> Ipynb.MimeBundle -> m UnembeddedMimeBundle
    go :: forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
Text -> MimeBundle -> m UnembeddedMimeBundle
go Text
attName = (Text -> ByteString -> FilePath)
-> FilePath -> MimeBundle -> m UnembeddedMimeBundle
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
(Text -> ByteString -> FilePath)
-> FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundleWith (Text -> Text -> ByteString -> FilePath
genFileName Text
attName) FilePath
subdir

    -- Include the attachment name in the bytes before hashing so that
    -- different attachments with the same content have different hashes.
    -- This is needed so that we know which attachment to reference when
    -- repacking the notebook.
    genFileName :: Text -> Text -> ByteString -> FilePath
    genFileName :: Text -> Text -> ByteString -> FilePath
genFileName Text
attName Text
mimeType ByteString
bytes = Text -> ByteString -> FilePath
binaryOutputFileName Text
mimeType (Text -> ByteString
Text.encodeUtf8 Text
attName ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
bytes)

unembedMimeBundleWith ::
  (MonadState [(FilePath, ByteString)] m) =>
  (Text -> ByteString -> FilePath) ->
  FilePath ->
  Ipynb.MimeBundle ->
  m UnembeddedMimeBundle
unembedMimeBundleWith :: forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
(Text -> ByteString -> FilePath)
-> FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundleWith Text -> ByteString -> FilePath
genFileName FilePath
subdir MimeBundle
mimeBundle = Map Text UnembeddedMimeData -> UnembeddedMimeBundle
UnembeddedMimeBundle (Map Text UnembeddedMimeData -> UnembeddedMimeBundle)
-> m (Map Text UnembeddedMimeData) -> m UnembeddedMimeBundle
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Text -> MimeData -> m UnembeddedMimeData)
-> Map Text MimeData -> m (Map Text UnembeddedMimeData)
forall (t :: * -> *) k a b.
Applicative t =>
(k -> a -> t b) -> Map k a -> t (Map k b)
Map.traverseWithKey Text -> MimeData -> m UnembeddedMimeData
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
Text -> MimeData -> m UnembeddedMimeData
go (MimeBundle -> Map Text MimeData
forall a b. Coercible a b => a -> b
coerce MimeBundle
mimeBundle)
  where
    go :: (MonadState [(FilePath, ByteString)] m) => Ipynb.MimeType -> Ipynb.MimeData -> m UnembeddedMimeData
    go :: forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
Text -> MimeData -> m UnembeddedMimeData
go Text
mimeType MimeData
mimeData = do
      let (UnembeddedMimeData
uMimeData, Maybe (FilePath, ByteString)
maybeExport) = (Text -> ByteString -> FilePath)
-> FilePath
-> Text
-> MimeData
-> (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeDataWith Text -> ByteString -> FilePath
genFileName FilePath
subdir Text
mimeType MimeData
mimeData
      ([(FilePath, ByteString)] -> [(FilePath, ByteString)]) -> m ()
forall s (m :: * -> *). MonadState s m => (s -> s) -> m ()
State.modify (([(FilePath, ByteString)] -> [(FilePath, ByteString)])
-> ((FilePath, ByteString)
    -> [(FilePath, ByteString)] -> [(FilePath, ByteString)])
-> Maybe (FilePath, ByteString)
-> [(FilePath, ByteString)]
-> [(FilePath, ByteString)]
forall b a. b -> (a -> b) -> Maybe a -> b
Maybe.maybe [(FilePath, ByteString)] -> [(FilePath, ByteString)]
forall a. a -> a
id (:) Maybe (FilePath, ByteString)
maybeExport)
      UnembeddedMimeData -> m UnembeddedMimeData
forall a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure UnembeddedMimeData
uMimeData

unembedMimeBundle :: (MonadState [(FilePath, ByteString)] m) => FilePath -> Ipynb.MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundle :: forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundle = (Text -> ByteString -> FilePath)
-> FilePath -> MimeBundle -> m UnembeddedMimeBundle
forall (m :: * -> *).
MonadState [(FilePath, ByteString)] m =>
(Text -> ByteString -> FilePath)
-> FilePath -> MimeBundle -> m UnembeddedMimeBundle
unembedMimeBundleWith Text -> ByteString -> FilePath
binaryOutputFileName

unembedMimeDataWith ::
  (Text -> ByteString -> FilePath) ->
  FilePath ->
  Ipynb.MimeType ->
  Ipynb.MimeData ->
  (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeDataWith :: (Text -> ByteString -> FilePath)
-> FilePath
-> Text
-> MimeData
-> (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeDataWith Text -> ByteString -> FilePath
genFileName FilePath
subdir Text
mimetype (Ipynb.BinaryData ByteString
bytes) =
  let filename :: FilePath
filename = Text -> ByteString -> FilePath
genFileName Text
mimetype ByteString
bytes
      filepath :: FilePath
filepath = FilePath
subdir FilePath -> FilePath -> FilePath
</> FilePath
filename
   in (FilePath -> UnembeddedMimeData
BinaryData FilePath
filepath, (FilePath, ByteString) -> Maybe (FilePath, ByteString)
forall a. a -> Maybe a
Just (FilePath
filepath, ByteString
bytes))
unembedMimeDataWith Text -> ByteString -> FilePath
_ FilePath
_ Text
_ (Ipynb.TextualData Text
text) = (Text -> UnembeddedMimeData
TextualData Text
text, Maybe (FilePath, ByteString)
forall a. Maybe a
Nothing)
unembedMimeDataWith Text -> ByteString -> FilePath
_ FilePath
_ Text
_ (Ipynb.JsonData Value
value) = (Value -> UnembeddedMimeData
JsonData Value
value, Maybe (FilePath, ByteString)
forall a. Maybe a
Nothing)

unembedMimeData ::
  FilePath ->
  Ipynb.MimeType ->
  Ipynb.MimeData ->
  (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeData :: FilePath
-> Text
-> MimeData
-> (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeData = (Text -> ByteString -> FilePath)
-> FilePath
-> Text
-> MimeData
-> (UnembeddedMimeData, Maybe (FilePath, ByteString))
unembedMimeDataWith Text -> ByteString -> FilePath
binaryOutputFileName

binaryOutputFileName :: Text -> ByteString -> FilePath
binaryOutputFileName :: Text -> ByteString -> FilePath
binaryOutputFileName Text
mimetype ByteString
bytes =
  Digest SHA256 -> FilePath
forall a. Digest a -> FilePath
digestToBase64 (SHA256 -> ByteString -> Digest SHA256
forall ba alg.
(ByteArrayAccess ba, HashAlgorithm alg) =>
alg -> ba -> Digest alg
Hash.hashWith SHA256
Hash.SHA256 ByteString
bytes)
    FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> case Text -> Maybe Text
extensionFromMimeType Text
mimetype of
      Maybe Text
Nothing -> FilePath
""
      Just Text
ext -> FilePath
"." FilePath -> FilePath -> FilePath
forall a. Semigroup a => a -> a -> a
<> Text -> FilePath
Text.unpack Text
ext
  where
    digestToBase64 :: Hash.Digest a -> String
    digestToBase64 :: forall a. Digest a -> FilePath
digestToBase64 Digest a
digest = ByteString -> FilePath
ByteString.unpack (ByteString -> FilePath)
-> (ByteString -> ByteString) -> ByteString -> FilePath
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
Base64.URL.encode (ByteString -> FilePath) -> ByteString -> FilePath
forall a b. (a -> b) -> a -> b
$ Digest a -> ByteString
forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
ByteArray.convert Digest a
digest

extensionFromMimeType :: Text -> Maybe Text
extensionFromMimeType :: Text -> Maybe Text
extensionFromMimeType Text
"text/plain" = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"txt"
extensionFromMimeType Text
"image/tiff" = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"tiff"
extensionFromMimeType Text
"image/jpeg" = Text -> Maybe Text
forall a. a -> Maybe a
Just Text
"jpg"
extensionFromMimeType Text
mt = do
  [Text]
exts <- ByteString -> Map ByteString [Text] -> Maybe [Text]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (Text -> ByteString
Text.encodeUtf8 Text
mt) Map ByteString [Text]
Mime.defaultExtensionMap
  [Text] -> Maybe Text
forall a. [a] -> Maybe a
Maybe.listToMaybe [Text]
exts