{-# LANGUAGE OverloadedStrings, ScopedTypeVariables #-}
{-# OPTIONS_GHC -fno-warn-unused-binds -fno-warn-name-shadowing #-}

-- | Description : Parsing messages received from IPython
--
-- This module is responsible for converting from low-level ByteStrings obtained from the 0MQ
-- sockets into Messages. The only exposed function is `parseMessage`, which should only be used in
-- the low-level 0MQ interface.
module IHaskell.IPython.Message.Parser (parseMessage) where

import           Control.Applicative ((<$>), (<*>))
import           Data.Aeson ((.:), (.:?), (.!=), decode, FromJSON, Result(..), Object, Value(..))
import           Data.Aeson.Types (Parser, parse, parseEither)
import           Data.ByteString hiding (unpack)
import qualified Data.ByteString.Lazy as Lazy
import           Data.HashMap.Strict as HM
import           Data.Maybe (fromMaybe)
import           Data.Text (unpack)
import           Debug.Trace
import           IHaskell.IPython.Types

type LByteString = Lazy.ByteString

-- --- External interface ----- | Parse a message from its ByteString components into a Message.
--   See https://jupyter-client.readthedocs.io/en/stable/messaging.html#the-wire-protocol
parseMessage :: [ByteString] -- ^ The list of identifiers sent with the message.
             -> ByteString   -- ^ The header data.
             -> ByteString   -- ^ The parent header, which is just "{}" if there is no header.
             -> ByteString   -- ^ The metadata map, also "{}" for an empty map.
             -> ByteString   -- ^ The message content.
             -> [ByteString] -- ^ Extra raw data buffer(s)
             -> Message      -- ^ A parsed message.
parseMessage idents headerData parentHeader metadata content buffers =
  let header = parseHeader idents headerData parentHeader metadata buffers
      messageType = mhMsgType header
      messageWithoutHeader = parser messageType $ Lazy.fromStrict content
  in messageWithoutHeader { header = header }

-- --- Module internals ----- | Parse a header from its ByteString components into a MessageHeader.
parseHeader :: [ByteString]  -- ^ The list of identifiers.
            -> ByteString    -- ^ The header data.
            -> ByteString    -- ^ The parent header, or "{}" for Nothing.
            -> ByteString    -- ^ The metadata, or "{}" for an empty map.
            -> [ByteString]  -- ^ Extra raw data buffer(s)
            -> MessageHeader -- The resulting message header.
parseHeader idents headerData parentHeader metadata buffers =
  MessageHeader idents parentResult metadataMap messageUUID sessionUUID username messageType buffers
  where
    -- Decode the header data and the parent header data into JSON objects. If the parent header data is
    -- absent, just have Nothing instead.
    Just result = decode $ Lazy.fromStrict headerData :: Maybe Object
    parentResult = if parentHeader == "{}"
                     then Nothing
                     else Just $ parseHeader idents parentHeader "{}" metadata []

    Success (messageType, username, messageUUID, sessionUUID) = flip parse result $ \obj -> do
      messType <- obj .: "msg_type"
      username <- obj .: "username"
      message <- obj .: "msg_id"
      session <- obj .: "session"
      return (messType, username, message, session)

    -- Get metadata as a simple map.
    Just metadataMap = fmap Metadata $ decode $ Lazy.fromStrict metadata

noHeader :: MessageHeader
noHeader = error "No header created"

parser :: MessageType            -- ^ The message type being parsed.
       -> LByteString -> Message -- ^ The parser that converts the body into a message. This message
                                 -- should have an undefined header.
parser KernelInfoRequestMessage = kernelInfoRequestParser
parser ExecuteInputMessage = executeInputParser
parser ExecuteRequestMessage = executeRequestParser
parser ExecuteReplyMessage = executeReplyParser
parser ExecuteErrorMessage = executeErrorParser
parser ExecuteResultMessage = executeResultParser
parser DisplayDataMessage = displayDataParser
parser IsCompleteRequestMessage = isCompleteRequestParser
parser CompleteRequestMessage = completeRequestParser
parser InspectRequestMessage = inspectRequestParser
parser ShutdownRequestMessage = shutdownRequestParser
parser InputReplyMessage = inputReplyParser
parser CommOpenMessage = commOpenParser
parser CommDataMessage = commDataParser
parser CommInfoRequestMessage = commInfoRequestParser
parser CommCloseMessage = commCloseParser
parser HistoryRequestMessage = historyRequestParser
parser StatusMessage = statusMessageParser
parser StreamMessage = streamMessageParser
parser InputMessage = inputMessageParser
parser OutputMessage = outputMessageParser
parser ClearOutputMessage = clearOutputMessageParser
parser other = error $ "Unknown message type " ++ show other

-- | Parse a kernel info request. A kernel info request has no auxiliary information, so ignore the
-- body.
kernelInfoRequestParser :: LByteString -> Message
kernelInfoRequestParser _ = KernelInfoRequest { header = noHeader }

-- | Parse a comm info request. A comm info request has no auxiliary information, so ignore the
-- body.
commInfoRequestParser :: LByteString -> Message
commInfoRequestParser _ = CommInfoRequest { header = noHeader }

-- | Parse an execute_input response. Fields used are:
executeInputParser :: LByteString -> Message
executeInputParser = requestParser $ \obj -> do
  code <- obj .: "code"
  executionCount <- obj .: "execution_count"
  return $ ExecuteInput noHeader code executionCount

-- | Parse an execute request. Fields used are:
--  1. "code": the code to execute.
--  2. "silent": whether to execute silently.
--  3. "store_history": whether to include this in history.
--  4. "allow_stdin": whether to allow reading from stdin for this code.
executeRequestParser :: LByteString -> Message
executeRequestParser content =
  let parser obj = do
                     let getOrElse a k = (fromMaybe a) <$> obj .:? k
                     code <- obj .: "code"
                     silent <- getOrElse False "silent"
                     storeHistory <- getOrElse (not silent) "store_history"
                     allowStdin <- obj .: "allow_stdin"

                     return (code, silent, storeHistory, allowStdin)
      Just decoded = decode content
      Success (code, silent, storeHistory, allowStdin) = parse parser decoded
  in ExecuteRequest
    { header = noHeader
    , getCode = code
    , getSilent = silent
    , getAllowStdin = allowStdin
    , getStoreHistory = storeHistory
    , getUserVariables = []
    , getUserExpressions = []
    }

-- | Parse an execute reply
executeReplyParser :: LByteString -> Message
executeReplyParser = requestParser $ \obj -> do
  status <- obj .: "status"
  executionCount <- obj .: "execution_count"
  return $ ExecuteReply noHeader status [] executionCount

-- | Parse an execute reply
executeErrorParser :: LByteString -> Message
executeErrorParser = requestParser $ \obj -> do
  -- executionCount <- obj .: "execution_count"
  traceback <- obj .: "traceback"
  ename <- obj .: "ename"
  evalue <- obj .: "evalue"
  return $ ExecuteError noHeader traceback ename evalue

makeDisplayDatas :: Object -> [DisplayData]
makeDisplayDatas dataDict = [DisplayData (read $ unpack mimeType) content | (mimeType, String content) <- HM.toList
                                                                                                            dataDict]

-- | Parse an execute result
executeResultParser :: LByteString -> Message
executeResultParser = requestParser $ \obj -> do
  executionCount <- obj .: "execution_count"
  dataDict :: Object <- obj .: "data"
  let displayDatas = makeDisplayDatas dataDict
  metadataDict <- obj .: "metadata"
  return $ ExecuteResult noHeader displayDatas metadataDict executionCount

-- | Parse a display data message
displayDataParser :: LByteString -> Message
displayDataParser = requestParser $ \obj -> do
  dataDict :: Object <- obj .: "data"
  let displayDatas = makeDisplayDatas dataDict
  return $ PublishDisplayData noHeader displayDatas Nothing

requestParser :: FromJSON a => (a -> Parser Message) -> LByteString -> Message
requestParser parser content =
  case parseEither parser decoded of
    Right parsed -> parsed
    Left err     -> trace ("Parse error: " ++ show err) SendNothing
  where
    Just decoded = decode content

historyRequestParser :: LByteString -> Message
historyRequestParser = requestParser $ \obj ->
  HistoryRequest noHeader <$> obj .: "output" <*> obj .: "raw" <*> historyAccessType obj
  where
    -- TODO: Implement full history access type parsing from message spec
    historyAccessType obj = do
      accessTypeStr <- obj .: "hist_access_type"
      return $
        case accessTypeStr of
          "range"  -> HistoryRange
          "tail"   -> HistoryTail
          "search" -> HistorySearch
          str      -> error $ "Unknown history access type: " ++ str

statusMessageParser :: LByteString -> Message
statusMessageParser = requestParser $ \obj -> do
  execution_state <- obj .: "execution_state"
  return $ PublishStatus noHeader execution_state

streamMessageParser :: LByteString -> Message
streamMessageParser = requestParser $ \obj -> do
  streamType <- obj .: "name"
  streamContent <- obj .: "text"
  return $ PublishStream noHeader streamType streamContent

inputMessageParser :: LByteString -> Message
inputMessageParser = requestParser $ \obj -> do
  code <- obj .: "code"
  executionCount <- obj .: "execution_count"
  return $ Input noHeader code executionCount

getDisplayDatas :: Maybe Object -> [DisplayData]
getDisplayDatas Nothing = []
getDisplayDatas (Just dataDict) = makeDisplayDatas dataDict

outputMessageParser :: LByteString -> Message
outputMessageParser = requestParser $ \obj -> do
  -- Handle both "data" and "text" keys
  maybeDataDict1 :: Maybe Object <- obj .:? "data"
  let displayDatas1 = getDisplayDatas maybeDataDict1

  maybeDataDict2 :: Maybe Object <- obj .:? "text"
  let displayDatas2 = getDisplayDatas maybeDataDict2

  executionCount <- obj .: "execution_count"
  return $ Output noHeader (displayDatas1 ++ displayDatas2) executionCount

clearOutputMessageParser :: LByteString -> Message
clearOutputMessageParser = requestParser $ \obj -> do
  wait <- obj .: "wait"
  return $ ClearOutput noHeader wait

isCompleteRequestParser :: LByteString -> Message
isCompleteRequestParser = requestParser $ \obj -> do
  code <- obj .: "code"
  return $ IsCompleteRequest noHeader code

completeRequestParser :: LByteString -> Message
completeRequestParser = requestParser $ \obj -> do
  code <- obj .: "code"
  pos <- obj .: "cursor_pos"
  return $ CompleteRequest noHeader code pos

inspectRequestParser :: LByteString -> Message
inspectRequestParser = requestParser $ \obj -> do
  code <- obj .: "code"
  pos <- obj .: "cursor_pos"
  dlevel <- obj .: "detail_level"
  return $ InspectRequest noHeader code pos dlevel

shutdownRequestParser :: LByteString -> Message
shutdownRequestParser = requestParser $ \obj -> do
  code <- obj .: "restart"
  return $ ShutdownRequest noHeader code

inputReplyParser :: LByteString -> Message
inputReplyParser = requestParser $ \obj -> do
  value <- obj .: "value"
  return $ InputReply noHeader value

commOpenParser :: LByteString -> Message
commOpenParser = requestParser $ \obj -> do
  uuid <- obj .: "comm_id"
  targetName <- obj .: "target_name"
  targetModule <- obj .:? "target_module" .!= ""
  value <- obj .: "data"
  return $ CommOpen noHeader targetName targetModule uuid value

commDataParser :: LByteString -> Message
commDataParser = requestParser $ \obj -> do
  uuid <- obj .: "comm_id"
  value <- obj .: "data"
  return $ CommData noHeader uuid value

commCloseParser :: LByteString -> Message
commCloseParser = requestParser $ \obj -> do
  uuid <- obj .: "comm_id"
  value <- obj .: "data"
  return $ CommClose noHeader uuid value