{-# LANGUAGE NoImplicitPrelude, OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE CPP #-}

module IHaskell.Convert.IpynbToLhs (ipynbToLhs) where

import           IHaskellPrelude
import qualified Data.Text.Lazy as LT
import qualified Data.ByteString.Lazy as LBS

import           Data.Aeson (decode, Object, Value(Array, Object, String))
import           Data.Vector (Vector)

import qualified Data.Text.Lazy.IO as LTIO
import qualified Data.Vector as V (map, mapM, toList)

import           IHaskell.Flags (LhsStyle(..))

#if MIN_VERSION_aeson(2,0,0)
import           Data.Aeson.KeyMap (lookup)
#else
import           Data.HashMap.Strict (lookup)
#endif

ipynbToLhs :: LhsStyle LText
           -> FilePath -- ^ the filename of an ipython notebook
           -> FilePath -- ^ the filename of the literate haskell to write
           -> IO ()
ipynbToLhs :: LhsStyle Text -> FilePath -> FilePath -> IO ()
ipynbToLhs LhsStyle Text
sty FilePath
from FilePath
to = do
  Just (Object
js :: Object) <- forall a. FromJSON a => ByteString -> Maybe a
decode forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO ByteString
LBS.readFile FilePath
from
  case forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cells" Object
js of
    Just (Array Array
cells) ->
      FilePath -> Text -> IO ()
LTIO.writeFile FilePath
to forall a b. (a -> b) -> a -> b
$ [Text] -> Text
LT.unlines forall a b. (a -> b) -> a -> b
$ forall a. Vector a -> [a]
V.toList forall a b. (a -> b) -> a -> b
$ forall a b. (a -> b) -> Vector a -> Vector b
V.map (\(Object Object
y) -> LhsStyle Text -> Object -> Text
convCell LhsStyle Text
sty Object
y) Array
cells
    Maybe Value
_ -> forall a. HasCallStack => FilePath -> a
error FilePath
"IHaskell.Convert.ipynbTolhs: json does not follow expected schema"

concatWithPrefix :: LT.Text -- ^ the prefix to add to every line
                 -> Vector Value          -- ^ a json array of text lines
                 -> Maybe LT.Text
concatWithPrefix :: Text -> Array -> Maybe Text
concatWithPrefix Text
p Array
arr = [Text] -> Text
LT.concat forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a b. (a -> b) -> [a] -> [b]
map (Text
p forall a. Semigroup a => a -> a -> a
<>) forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. Vector a -> [a]
V.toList forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM Value -> Maybe Text
toStr Array
arr

toStr :: Value -> Maybe LT.Text
toStr :: Value -> Maybe Text
toStr (String Text
x) = forall a. a -> Maybe a
Just (Text -> Text
LT.fromStrict Text
x)
toStr Value
_ = forall a. Maybe a
Nothing

-- | @convCell sty cell@ converts a single cell in JSON into text suitable for the type of lhs file
-- described by the @sty@
convCell :: LhsStyle LT.Text -> Object -> LT.Text
convCell :: LhsStyle Text -> Object -> Text
convCell LhsStyle Text
_sty Object
object
  | Just (String Text
"markdown") <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cell_type" Object
object,
    Just (Array Array
xs) <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"source" Object
object,
    ~(Just Text
s) <- Text -> Array -> Maybe Text
concatWithPrefix Text
"" Array
xs
  = Text
s
convCell LhsStyle Text
sty Object
object
  | Just (String Text
"code") <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cell_type" Object
object,
    Just (Array Array
a) <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"source" Object
object,
    Just (Array Array
o) <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"outputs" Object
object,
    ~(Just Text
i) <- Text -> Array -> Maybe Text
concatWithPrefix (forall string. LhsStyle string -> string
lhsCodePrefix LhsStyle Text
sty) Array
a,
    Text
o2 <- forall a. a -> Maybe a -> a
fromMaybe forall a. Monoid a => a
mempty (LhsStyle Text -> Array -> Maybe Text
convOutputs LhsStyle Text
sty Array
o)
  = Text
"\n" forall a. Semigroup a => a -> a -> a
<>
    forall string. LhsStyle string -> string
lhsBeginCode LhsStyle Text
sty forall a. Semigroup a => a -> a -> a
<> Text
i forall a. Semigroup a => a -> a -> a
<> forall string. LhsStyle string -> string
lhsEndCode LhsStyle Text
sty forall a. Semigroup a => a -> a -> a
<> Text
"\n" forall a. Semigroup a => a -> a -> a
<> Text
o2 forall a. Semigroup a => a -> a -> a
<> Text
"\n"
convCell LhsStyle Text
_ Object
_ = Text
"IHaskell.Convert.convCell: unknown cell"

convOutputs :: LhsStyle LT.Text
            -> Vector Value -- ^ JSON array of output lines containing text or markup
            -> Maybe LT.Text
convOutputs :: LhsStyle Text -> Array -> Maybe Text
convOutputs LhsStyle Text
sty Array
array = do
  Vector Text
outputLines <- forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (Text -> Value -> Maybe Text
getTexts (forall string. LhsStyle string -> string
lhsOutputPrefix LhsStyle Text
sty)) Array
array
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall string. LhsStyle string -> string
lhsBeginOutput LhsStyle Text
sty forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
LT.concat (forall a. Vector a -> [a]
V.toList Vector Text
outputLines) forall a. Semigroup a => a -> a -> a
<> forall string. LhsStyle string -> string
lhsEndOutput LhsStyle Text
sty

getTexts :: LT.Text -> Value -> Maybe LT.Text
getTexts :: Text -> Value -> Maybe Text
getTexts Text
p (Object Object
object)
  | Just (Array Array
text) <- forall v. Key -> KeyMap v -> Maybe v
lookup Key
"text" Object
object = Text -> Array -> Maybe Text
concatWithPrefix Text
p Array
text
getTexts Text
_ Value
_ = forall a. Maybe a
Nothing