{-# 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) <- ByteString -> Maybe Object
forall a. FromJSON a => ByteString -> Maybe a
decode (ByteString -> Maybe Object) -> IO ByteString -> IO (Maybe Object)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FilePath -> IO ByteString
LBS.readFile FilePath
from
  case Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cells" Object
js of
    Just (Array Array
cells) ->
      FilePath -> Text -> IO ()
LTIO.writeFile FilePath
to (Text -> IO ()) -> Text -> IO ()
forall a b. (a -> b) -> a -> b
$ [Text] -> Text
LT.unlines ([Text] -> Text) -> [Text] -> Text
forall a b. (a -> b) -> a -> b
$ Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList (Vector Text -> [Text]) -> Vector Text -> [Text]
forall a b. (a -> b) -> a -> b
$ (Value -> Text) -> Array -> Vector Text
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
_ -> FilePath -> IO ()
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 ([Text] -> Text) -> (Vector Text -> [Text]) -> Vector Text -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> Text) -> [Text] -> [Text]
forall a b. (a -> b) -> [a] -> [b]
map (Text
p Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>) ([Text] -> [Text])
-> (Vector Text -> [Text]) -> Vector Text -> [Text]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList (Vector Text -> Text) -> Maybe (Vector Text) -> Maybe Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Value -> Maybe Text) -> Array -> Maybe (Vector Text)
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) = Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Text
LT.fromStrict Text
x)
toStr Value
_ = Maybe Text
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") <- Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cell_type" Object
object,
    Just (Array Array
xs) <- Key -> Object -> Maybe Value
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") <- Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"cell_type" Object
object,
    Just (Array Array
a) <- Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"source" Object
object,
    Just (Array Array
o) <- Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"outputs" Object
object,
    ~(Just Text
i) <- Text -> Array -> Maybe Text
concatWithPrefix (LhsStyle Text -> Text
forall string. LhsStyle string -> string
lhsCodePrefix LhsStyle Text
sty) Array
a,
    Text
o2 <- Text -> Maybe Text -> Text
forall a. a -> Maybe a -> a
fromMaybe Text
forall a. Monoid a => a
mempty (LhsStyle Text -> Array -> Maybe Text
convOutputs LhsStyle Text
sty Array
o)
  = Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<>
    LhsStyle Text -> Text
forall string. LhsStyle string -> string
lhsBeginCode LhsStyle Text
sty Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
i Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LhsStyle Text -> Text
forall string. LhsStyle string -> string
lhsEndCode LhsStyle Text
sty Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
"\n" Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> Text
o2 Text -> Text -> Text
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 <- (Value -> Maybe Text) -> Array -> Maybe (Vector Text)
forall (m :: * -> *) a b.
Monad m =>
(a -> m b) -> Vector a -> m (Vector b)
V.mapM (Text -> Value -> Maybe Text
getTexts (LhsStyle Text -> Text
forall string. LhsStyle string -> string
lhsOutputPrefix LhsStyle Text
sty)) Array
array
  Text -> Maybe Text
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return (Text -> Maybe Text) -> Text -> Maybe Text
forall a b. (a -> b) -> a -> b
$ LhsStyle Text -> Text
forall string. LhsStyle string -> string
lhsBeginOutput LhsStyle Text
sty Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> [Text] -> Text
LT.concat (Vector Text -> [Text]
forall a. Vector a -> [a]
V.toList Vector Text
outputLines) Text -> Text -> Text
forall a. Semigroup a => a -> a -> a
<> LhsStyle Text -> Text
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) <- Key -> Object -> Maybe Value
forall v. Key -> KeyMap v -> Maybe v
lookup Key
"text" Object
object = Text -> Array -> Maybe Text
concatWithPrefix Text
p Array
text
getTexts Text
_ Value
_ = Maybe Text
forall a. Maybe a
Nothing