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

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           Data.HashMap.Strict (lookup)

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

import           IHaskell.Flags (LhsStyle(..))

ipynbToLhs :: LhsStyle LText
           -> FilePath -- ^ the filename of an ipython notebook
           -> FilePath -- ^ the filename of the literate haskell to write
           -> IO ()
ipynbToLhs sty from to = do
  Just (js :: Object) <- decode <$> LBS.readFile from
  case lookup "cells" js of
    Just (Array cells) ->
      LTIO.writeFile to $ LT.unlines $ V.toList $ V.map (\(Object y) -> convCell sty y) cells
    _ -> error "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 p arr = LT.concat . map (p <>) . V.toList <$> V.mapM toStr arr

toStr :: Value -> Maybe LT.Text
toStr (String x) = Just (LT.fromStrict x)
toStr _ = 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 _sty object
  | Just (String "markdown") <- lookup "cell_type" object,
    Just (Array xs) <- lookup "source" object,
    ~(Just s) <- concatWithPrefix "" xs
  = s
convCell sty object
  | Just (String "code") <- lookup "cell_type" object,
    Just (Array a) <- lookup "source" object,
    Just (Array o) <- lookup "outputs" object,
    ~(Just i) <- concatWithPrefix (lhsCodePrefix sty) a,
    o2 <- fromMaybe mempty (convOutputs sty o)
  = "\n" <>
    lhsBeginCode sty <> i <> lhsEndCode sty <> "\n" <> o2 <> "\n"
convCell _ _ = "IHaskell.Convert.convCell: unknown cell"

convOutputs :: LhsStyle LT.Text
            -> Vector Value -- ^ JSON array of output lines containing text or markup
            -> Maybe LT.Text
convOutputs sty array = do
  outputLines <- V.mapM (getTexts (lhsOutputPrefix sty)) array
  return $ lhsBeginOutput sty <> LT.concat (V.toList outputLines) <> lhsEndOutput sty

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