{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE FlexibleContexts  #-}
{-# LANGUAGE ScopedTypeVariables #-}
{- |
   Module      : Text.Pandoc.Readers.Ipynb
   Copyright   : Copyright (C) 2019-2024 John MacFarlane
   License     : GNU GPL, version 2 or above

   Maintainer  : John MacFarlane <jgm@berkeley.edu>
   Stability   : alpha
   Portability : portable

Ipynb (Jupyter notebook JSON format) reader for pandoc.
-}
module Text.Pandoc.Readers.Ipynb ( readIpynb )
where
import Crypto.Hash (hashWith, SHA1(SHA1))
import Data.Char (isDigit)
import Data.Maybe (fromMaybe)
import Text.Pandoc.Options
import Control.Applicative ((<|>))
import qualified Data.Scientific as Scientific
import qualified Text.Pandoc.Builder as B
import Text.Pandoc.Logging
import Text.Pandoc.Definition
import Data.Ipynb as Ipynb
import Text.Pandoc.Class.PandocMonad
import Text.Pandoc.MIME (extensionFromMimeType)
import Text.Pandoc.Shared (tshow)
import Text.Pandoc.UTF8
import Text.Pandoc.Walk (walk)
import Text.Pandoc.Error
import Data.Text (Text)
import qualified Data.Map as M
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.ByteString.Lazy as BL
import Data.Aeson as Aeson
import Control.Monad.Except (throwError)
import Text.Pandoc.Readers.Markdown (readMarkdown)
import qualified Text.Pandoc.UTF8 as UTF8
import Text.Pandoc.Sources (ToSources(..), sourcesToText)

readIpynb :: (PandocMonad m, ToSources a)
          => ReaderOptions -> a -> m Pandoc
readIpynb opts x = do
  let src = BL.fromStrict . TE.encodeUtf8 . sourcesToText $ toSources x
  case eitherDecode src of
    Right (notebook4 :: Notebook NbV4) -> notebookToPandoc opts notebook4
    Left _ ->
      case eitherDecode src of
        Right (notebook3 :: Notebook NbV3) -> notebookToPandoc opts notebook3
        Left err -> throwError $ PandocIpynbDecodingError $ T.pack err

notebookToPandoc :: PandocMonad m
                 => ReaderOptions -> Notebook a -> m Pandoc
notebookToPandoc opts notebook = do
  let cells = notebookCells notebook
  let (fmt,fmtminor) = notebookFormat notebook
  let m = M.insert "nbformat" (MetaString $ tshow fmt) $
          M.insert "nbformat_minor" (MetaString $ tshow fmtminor) $
          jsonMetaToMeta (notebookMetadata notebook)
  let lang = case M.lookup "kernelspec" m of
                   Just (MetaMap ks) ->
                      case M.lookup "language" ks of
                         Just (MetaString l) -> l
                         _ -> "python"
                   _ -> "python"
  bs <- mconcat <$> mapM (cellToBlocks opts lang) cells
  let Pandoc _ blocks = B.doc bs
  return $ Pandoc (Meta $ M.insert "jupyter" (MetaMap m) mempty) blocks

cellToBlocks :: PandocMonad m
             => ReaderOptions -> Text -> Ipynb.Cell a -> m B.Blocks
cellToBlocks opts lang c = do
  let Source ts = cellSource c
  let source = mconcat ts
  let kvs = jsonMetaToPairs (cellMetadata c)
  let attachments = case cellAttachments c of
                      Nothing -> mempty
                      Just (MimeAttachments m) -> M.toList m
  let ident = fromMaybe mempty $ cellId c
  mapM_ (addAttachment (cellId c)) attachments
  case cellType c of
    Ipynb.Markdown -> do
      bs <- if isEnabled Ext_raw_markdown opts
               then return [RawBlock (Format "markdown") source]
               else do
                 Pandoc _ bs <- walk (fixImage (cellId c)) <$> readMarkdown opts source
                 return bs
      return $ B.divWith (ident,["cell","markdown"],kvs)
             $ B.fromList bs
    Ipynb.Heading lev -> do
      Pandoc _ bs <- readMarkdown opts
        (T.replicate lev "#" <> " " <> source)
      return $ B.divWith (ident,["cell","markdown"],kvs)
             $ B.fromList bs
    Ipynb.Raw -> do
      -- we use ipynb to indicate no format given (a wildcard in nbformat)
      let format = fromMaybe "ipynb" $ lookup "raw_mimetype" kvs <|> lookup "format" kvs
      let format' =
            case format of
              "text/html"             -> "html"
              "slides"                -> "html"
              "text/latex"            -> "latex"
              "application/pdf"       -> "latex"
              "pdf"                   -> "latex"
              "text/markdown"         -> "markdown"
              "text/x-rst"            -> "rst"
              "text/restructuredtext" -> "rst"
              "text/asciidoc"         -> "asciidoc"
              _                       -> format
      return $ B.divWith (ident,["cell","raw"],kvs)
             $ B.rawBlock format' source
    Ipynb.Code{ codeOutputs = outputs, codeExecutionCount = ec } -> do
      outputBlocks <- mconcat <$> mapM outputToBlock outputs
      let kvs' = maybe kvs (\x -> ("execution_count", tshow x):kvs) ec
      return $ B.divWith (ident,["cell","code"],kvs') $
        B.codeBlockWith ("",[lang],[]) source
        <> outputBlocks

-- Remove attachment: prefix from images...
fixImage :: Maybe Text -> Inline -> Inline
fixImage mbident (Image attr lab (src,tit))
  | "attachment:" `T.isPrefixOf` src =
     let src' = T.drop 11 src
         qualifiedSrc = maybe src' (<> ("-" <> src')) mbident
      in Image attr lab (qualifiedSrc, tit)
fixImage _ x = x

addAttachment :: PandocMonad m => Maybe Text -> (Text, MimeBundle) -> m ()
addAttachment mbident (fname, mimeBundle) = do
  let fp = T.unpack $ maybe fname (<> ("-" <> fname)) mbident
  case M.toList (unMimeBundle mimeBundle) of
    (mimeType, BinaryData bs):_ ->
      insertMedia fp (Just mimeType) (BL.fromStrict bs)
    (mimeType, TextualData t):_ ->
      insertMedia fp (Just mimeType)
          (BL.fromStrict $ TE.encodeUtf8 t)
    (mimeType, JsonData v):_ ->
      insertMedia fp (Just mimeType) (encode v)
    [] -> report $ CouldNotFetchResource fname "no attachment"

outputToBlock :: PandocMonad m => Output a -> m B.Blocks
outputToBlock Stream{ streamName = sName,
                      streamText = Source text } =
  return $ B.divWith ("",["output","stream",sName],[])
         $ B.codeBlock $ T.concat text
outputToBlock DisplayData{ displayData = data',
                            displayMetadata = metadata' } =
  B.divWith ("",["output", "display_data"],[]) <$>
    handleData metadata' data'
outputToBlock ExecuteResult{ executeCount = ec,
                              executeData = data',
                              executeMetadata = metadata' } =
  B.divWith ("",["output", "execute_result"],[("execution_count",tshow ec)])
    <$> handleData metadata' data'
outputToBlock Err{ errName = ename,
                   errValue = evalue,
                   errTraceback = traceback } =
  return $ B.divWith ("",["output","error"],
                         [("ename",ename),
                          ("evalue",evalue)])
         $ B.codeBlock $ T.unlines traceback

-- We want to display the richest output possible given
-- the output format.
handleData :: PandocMonad m
           => JSONMeta -> MimeBundle -> m B.Blocks
handleData (JSONMeta metadata) (MimeBundle mb) =
  mconcat <$> mapM dataBlock (M.toList mb)

  where

    dataBlock :: PandocMonad m => (MimeType, MimeData) -> m B.Blocks
    dataBlock (mt, d)
     | "image/" `T.isPrefixOf` mt || mt == "application/pdf"
      = do
      -- normally metadata maps from mime types to key-value map;
      -- but not always...
      let meta = case M.lookup mt metadata of
                   Just v@Object{} ->
                     case fromJSON v of
                       Success m' -> m'
                       Error _   -> mempty
                   _ -> mempty
      let metaPairs = jsonMetaToPairs meta
      let bs = case d of
                 BinaryData bs' -> bs'
                 TextualData t  -> UTF8.fromText t
                 JsonData v     -> BL.toStrict $ encode v
      -- SHA1 hash for filename
      let fname = T.pack (show (hashWith SHA1 bs)) <>
            case extensionFromMimeType mt of
              Nothing  -> ""
              Just ext -> "." <> ext
      insertMedia (T.unpack fname) (Just mt) (BL.fromStrict bs)
      return $ B.para $ B.imageWith ("",[],metaPairs) fname "" mempty

    dataBlock ("text/html", TextualData t)
      = return $ B.rawBlock "html" t

    dataBlock ("text/latex", TextualData t)
      = return $ B.rawBlock "latex" t

    dataBlock ("text/markdown", TextualData t)
      = return $ B.rawBlock "markdown" t

    dataBlock ("text/plain", TextualData t) =
      return $ B.codeBlock t

    dataBlock (_, JsonData v) =
      return $ B.codeBlockWith ("",["json"],[]) $ T.pack $ toStringLazy $ encode v

    dataBlock _ = return mempty

jsonMetaToMeta :: JSONMeta -> M.Map Text MetaValue
jsonMetaToMeta (JSONMeta m) = M.map valueToMetaValue m
  where
    valueToMetaValue :: Value -> MetaValue
    valueToMetaValue x@Object{} =
      case fromJSON x of
        Error s -> MetaString $ T.pack s
        Success jm' -> MetaMap $ jsonMetaToMeta jm'
    valueToMetaValue x@Array{} =
      case fromJSON x of
        Error s -> MetaString $ T.pack s
        Success xs -> MetaList $ map valueToMetaValue xs
    valueToMetaValue (Bool b) = MetaBool b
    valueToMetaValue (String t) = MetaString t
    valueToMetaValue (Number n)
      | Scientific.isInteger n = MetaString (tshow (floor n :: Integer))
      | otherwise              = MetaString (tshow n)
    valueToMetaValue Aeson.Null = MetaString ""

jsonMetaToPairs :: JSONMeta -> [(Text, Text)]
jsonMetaToPairs (JSONMeta m) = M.toList . M.map
  (\case
      String t
        | not (T.all isDigit t)
        , t /= "true"
        , t /= "false"
                 -> t
      x          -> T.pack $ UTF8.toStringLazy $ Aeson.encode x) $ m
