{-# language NoImplicitPrelude, DoAndIfThenElse, OverloadedStrings, ExtendedDefaultRules, CPP #-}
module IHaskell.Eval.Widgets (
    widgetSendOpen,
    widgetSendView,
    widgetSendUpdate,
    widgetSendCustom,
    widgetSendClose,
    widgetSendValue,
    widgetPublishDisplay,
    widgetClearOutput,
    relayWidgetMessages,
    widgetHandler,
    ) where

import           IHaskellPrelude

import           Control.Concurrent.STM (atomically)
import           Control.Concurrent.STM.TChan
import           Data.Aeson
import           Data.ByteString.Base64 as B64 (decodeLenient)
import qualified Data.Map as Map
import           Data.Text.Encoding (encodeUtf8)

import           Data.Foldable (foldl)
import           System.IO.Unsafe (unsafePerformIO)

import           IHaskell.Display
import           IHaskell.Eval.Util (unfoldM)
import           IHaskell.IPython.Types (showMessageType)
import           IHaskell.Types

#if MIN_VERSION_aeson(2,0,0)
import qualified Data.Aeson.KeyMap as KM (lookup,insert,delete)
import qualified Data.Aeson.Key    as Key
#else
import qualified Data.HashMap.Strict as HM (lookup,insert,delete)
#endif

-- All comm_open messages go here
widgetMessages :: TChan WidgetMsg
{-# NOINLINE widgetMessages #-}
widgetMessages :: TChan WidgetMsg
widgetMessages = forall a. IO a -> a
unsafePerformIO forall a. IO (TChan a)
newTChanIO

-- | Return all pending comm_close messages
relayWidgetMessages :: IO [WidgetMsg]
relayWidgetMessages :: IO [WidgetMsg]
relayWidgetMessages = forall a. TChan a -> IO [a]
relayMessages TChan WidgetMsg
widgetMessages

-- | Extract all messages from a TChan and wrap them in a list
relayMessages :: TChan a -> IO [a]
relayMessages :: forall a. TChan a -> IO [a]
relayMessages = forall a. IO (Maybe a) -> IO [a]
unfoldM forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. STM a -> IO a
atomically forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TChan a -> STM (Maybe a)
tryReadTChan

-- | Write a widget message to the chan
queue :: WidgetMsg -> IO ()
queue :: WidgetMsg -> IO ()
queue = forall a. STM a -> IO a
atomically forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. TChan a -> a -> STM ()
writeTChan TChan WidgetMsg
widgetMessages

-- | Send a message
widgetSend :: IHaskellWidget a
           => (Widget -> Value -> WidgetMsg)
           -> a -> Value -> IO ()
widgetSend :: forall a.
IHaskellWidget a =>
(Widget -> Value -> WidgetMsg) -> a -> Value -> IO ()
widgetSend Widget -> Value -> WidgetMsg
mtype a
widget Value
value = WidgetMsg -> IO ()
queue forall a b. (a -> b) -> a -> b
$ Widget -> Value -> WidgetMsg
mtype (forall a. IHaskellWidget a => a -> Widget
Widget a
widget) Value
value

-- | Send a message to open a comm
widgetSendOpen :: IHaskellWidget a => a -> Value -> IO ()
widgetSendOpen :: forall a. IHaskellWidget a => a -> Value -> IO ()
widgetSendOpen = forall a.
IHaskellWidget a =>
(Widget -> Value -> WidgetMsg) -> a -> Value -> IO ()
widgetSend Widget -> Value -> WidgetMsg
Open

-- | Send a state update message
widgetSendUpdate :: IHaskellWidget a => a -> Value -> IO ()
widgetSendUpdate :: forall a. IHaskellWidget a => a -> Value -> IO ()
widgetSendUpdate = forall a.
IHaskellWidget a =>
(Widget -> Value -> WidgetMsg) -> a -> Value -> IO ()
widgetSend Widget -> Value -> WidgetMsg
Update

-- | Send a [method .= display] comm_msg
widgetSendView :: IHaskellWidget a => a -> IO ()
widgetSendView :: forall a. IHaskellWidget a => a -> IO ()
widgetSendView = WidgetMsg -> IO ()
queue forall b c a. (b -> c) -> (a -> b) -> a -> c
. Widget -> WidgetMsg
View forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. IHaskellWidget a => a -> Widget
Widget

-- | Send a comm_close
widgetSendClose :: IHaskellWidget a => a -> Value -> IO ()
widgetSendClose :: forall a. IHaskellWidget a => a -> Value -> IO ()
widgetSendClose = forall a.
IHaskellWidget a =>
(Widget -> Value -> WidgetMsg) -> a -> Value -> IO ()
widgetSend Widget -> Value -> WidgetMsg
Close

-- | Send a [method .= custom, content .= value] comm_msg
widgetSendCustom :: IHaskellWidget a => a -> Value -> IO ()
widgetSendCustom :: forall a. IHaskellWidget a => a -> Value -> IO ()
widgetSendCustom = forall a.
IHaskellWidget a =>
(Widget -> Value -> WidgetMsg) -> a -> Value -> IO ()
widgetSend Widget -> Value -> WidgetMsg
Custom

-- | Send a custom Value
widgetSendValue :: IHaskellWidget a => a -> Value -> IO ()
widgetSendValue :: forall a. IHaskellWidget a => a -> Value -> IO ()
widgetSendValue a
widget = WidgetMsg -> IO ()
queue forall b c a. (b -> c) -> (a -> b) -> a -> c
. Widget -> Value -> WidgetMsg
JSONValue (forall a. IHaskellWidget a => a -> Widget
Widget a
widget)

-- | Send a `display_data` message as a [method .= custom] message
widgetPublishDisplay :: (IHaskellWidget a, IHaskellDisplay b) => a -> b -> IO ()
widgetPublishDisplay :: forall a b.
(IHaskellWidget a, IHaskellDisplay b) =>
a -> b -> IO ()
widgetPublishDisplay a
widget b
disp = forall a. IHaskellDisplay a => a -> IO Display
display b
disp forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= WidgetMsg -> IO ()
queue forall b c a. (b -> c) -> (a -> b) -> a -> c
. Widget -> Display -> WidgetMsg
DispMsg (forall a. IHaskellWidget a => a -> Widget
Widget a
widget)

-- | Send a `clear_output` message
widgetClearOutput :: Bool -> IO ()
widgetClearOutput :: Bool -> IO ()
widgetClearOutput Bool
w = WidgetMsg -> IO ()
queue forall a b. (a -> b) -> a -> b
$ Bool -> WidgetMsg
ClrOutput Bool
w

-- | Handle a single widget message. Takes necessary actions according to the message type, such as
-- opening comms, storing and updating widget representation in the kernel state etc.
handleMessage :: (Message -> IO ())
              -> MessageHeader
              -> KernelState
              -> WidgetMsg
              -> IO KernelState
handleMessage :: (Message -> IO ())
-> MessageHeader -> KernelState -> WidgetMsg -> IO KernelState
handleMessage Message -> IO ()
send MessageHeader
replyHeader KernelState
state WidgetMsg
msg = do
  case WidgetMsg
msg of
    Open Widget
widget Value
value -> do
      let target_name :: String
target_name = forall a. IHaskellWidget a => a -> String
targetName Widget
widget
          target_module :: String
target_module = forall a. IHaskellWidget a => a -> String
targetModule Widget
widget
          uuid :: UUID
uuid = forall a. IHaskellWidget a => a -> UUID
getCommUUID Widget
widget
          present :: Bool
present = forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup UUID
uuid Map UUID Widget
oldComms

          newComms :: Map UUID Widget
newComms = forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert UUID
uuid Widget
widget Map UUID Widget
oldComms
          newState :: KernelState
newState = KernelState
state { openComms :: Map UUID Widget
openComms = Map UUID Widget
newComms }

          (Value
newvalue,[ByteString]
buffers,[BufferPath]
bp) = Value -> [BufferPath] -> (Value, [ByteString], [BufferPath])
processBPs Value
value forall a b. (a -> b) -> a -> b
$ forall a. IHaskellWidget a => a -> [BufferPath]
getBufferPaths Widget
widget
          applyBuffers :: MessageHeader -> MessageHeader
applyBuffers MessageHeader
x = MessageHeader
x {mhBuffers :: [ByteString]
mhBuffers = [ByteString]
buffers}
          content :: Value
content = [Pair] -> Value
object [ Key
"state" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Value
newvalue, Key
"buffer_paths" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= [BufferPath]
bp ]

          communicate :: Value -> IO ()
communicate Value
val = do
            MessageHeader
head <- MessageHeader -> MessageHeader
applyBuffers forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
CommDataMessage
            Message -> IO ()
send forall a b. (a -> b) -> a -> b
$ MessageHeader -> UUID -> Value -> Message
CommData MessageHeader
head UUID
uuid Value
val

      -- If the widget is present, don't open it again.
      if Bool
present
        then forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
state
        else do
          -- Send the comm open, with the initial state
          MessageHeader
hdr <- MessageHeader -> MessageHeader
applyBuffers forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
CommOpenMessage
          let hdrV :: MessageHeader
hdrV = MessageHeader -> String -> MessageHeader
setVersion MessageHeader
hdr String
"2.0.0" -- Widget Messaging Protocol Version
          Message -> IO ()
send forall a b. (a -> b) -> a -> b
$ MessageHeader -> String -> String -> UUID -> Value -> Message
CommOpen MessageHeader
hdrV String
target_name String
target_module UUID
uuid Value
content

          -- Send anything else the widget requires.
          forall a. IHaskellWidget a => a -> (Value -> IO ()) -> IO ()
open Widget
widget Value -> IO ()
communicate

          -- Store the widget in the kernelState
          forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
newState

    Close Widget
widget Value
value -> do
      let uuid :: UUID
uuid = forall a. IHaskellWidget a => a -> UUID
getCommUUID Widget
widget
          present :: Bool
present = forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup UUID
uuid Map UUID Widget
oldComms

          newComms :: Map UUID Widget
newComms = forall k a. Ord k => k -> Map k a -> Map k a
Map.delete UUID
uuid forall a b. (a -> b) -> a -> b
$ KernelState -> Map UUID Widget
openComms KernelState
state
          newState :: KernelState
newState = KernelState
state { openComms :: Map UUID Widget
openComms = Map UUID Widget
newComms }

      -- If the widget is not present in the state, we don't close it.
      if Bool
present
        then do
          MessageHeader
hdr <- MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
CommCloseMessage
          Message -> IO ()
send forall a b. (a -> b) -> a -> b
$ MessageHeader -> UUID -> Value -> Message
CommClose MessageHeader
hdr UUID
uuid Value
value
          forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
newState
        else forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
state

    View Widget
widget -> forall {p}. IHaskellWidget p => p -> Value -> IO KernelState
sendMessage Widget
widget (forall a. ToJSON a => a -> Value
toJSON WidgetMethod
DisplayWidget)

    Update Widget
widget Value
value -> do
      let (Value
newvalue,[ByteString]
buffers,[BufferPath]
bp) = Value -> [BufferPath] -> (Value, [ByteString], [BufferPath])
processBPs Value
value forall a b. (a -> b) -> a -> b
$ forall a. IHaskellWidget a => a -> [BufferPath]
getBufferPaths Widget
widget
      forall {p}.
IHaskellWidget p =>
p -> Value -> (MessageHeader -> MessageHeader) -> IO KernelState
sendMessageHdr Widget
widget (forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ Value -> [BufferPath] -> WidgetMethod
UpdateState Value
newvalue [BufferPath]
bp) (\MessageHeader
h->MessageHeader
h {mhBuffers :: [ByteString]
mhBuffers=[ByteString]
buffers})

    Custom Widget
widget Value
value -> forall {p}. IHaskellWidget p => p -> Value -> IO KernelState
sendMessage Widget
widget (forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ Value -> WidgetMethod
CustomContent Value
value)

    JSONValue Widget
widget Value
value -> forall {p}. IHaskellWidget p => p -> Value -> IO KernelState
sendMessage Widget
widget Value
value

    DispMsg Widget
widget Display
disp -> do
      MessageHeader
dispHeader <- MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
DisplayDataMessage
      let dmsg :: WidgetDisplay
dmsg = MessageHeader -> [DisplayData] -> WidgetDisplay
WidgetDisplay MessageHeader
dispHeader forall a b. (a -> b) -> a -> b
$ Display -> [DisplayData]
unwrap Display
disp
      forall {p}. IHaskellWidget p => p -> Value -> IO KernelState
sendMessage Widget
widget (forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ Value -> WidgetMethod
CustomContent forall a b. (a -> b) -> a -> b
$ forall a. ToJSON a => a -> Value
toJSON WidgetDisplay
dmsg)

    ClrOutput Bool
w -> do
      MessageHeader
hdr <- MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
ClearOutputMessage
      Message -> IO ()
send forall a b. (a -> b) -> a -> b
$ MessageHeader -> Bool -> Message
ClearOutput MessageHeader
hdr Bool
w
      forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
state

  where
    oldComms :: Map UUID Widget
oldComms = KernelState -> Map UUID Widget
openComms KernelState
state
    sendMessage :: p -> Value -> IO KernelState
sendMessage p
widget Value
value = forall {p}.
IHaskellWidget p =>
p -> Value -> (MessageHeader -> MessageHeader) -> IO KernelState
sendMessageHdr p
widget Value
value forall a. a -> a
id
    sendMessageHdr :: p -> Value -> (MessageHeader -> MessageHeader) -> IO KernelState
sendMessageHdr p
widget Value
value MessageHeader -> MessageHeader
hdrf = do
      let uuid :: UUID
uuid = forall a. IHaskellWidget a => a -> UUID
getCommUUID p
widget
          present :: Bool
present = forall a. Maybe a -> Bool
isJust forall a b. (a -> b) -> a -> b
$ forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup UUID
uuid Map UUID Widget
oldComms

      -- If the widget is present, we send an update message on its comm.
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
when Bool
present forall a b. (a -> b) -> a -> b
$ do
        MessageHeader
hdr <- MessageHeader -> MessageHeader
hdrf forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> MessageHeader -> MessageType -> IO MessageHeader
dupHeader MessageHeader
replyHeader MessageType
CommDataMessage
        Message -> IO ()
send forall a b. (a -> b) -> a -> b
$ MessageHeader -> UUID -> Value -> Message
CommData MessageHeader
hdr UUID
uuid Value
value
      forall (m :: * -> *) a. Monad m => a -> m a
return KernelState
state

    unwrap :: Display -> [DisplayData]
    unwrap :: Display -> [DisplayData]
unwrap (ManyDisplay [Display]
ds) = forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap Display -> [DisplayData]
unwrap [Display]
ds
    unwrap (Display [DisplayData]
ddatas) = [DisplayData]
ddatas

    -- Removes the values that are buffers and puts them in the third value of the tuple
    -- The returned bufferpaths are the bufferpaths used
    processBPs :: Value -> [BufferPath] -> (Value, [ByteString], [BufferPath])
    -- Searching if the BufferPath key is in the Object is O(log n) or O(1) depending on implementation
    -- For this reason we fold on the bufferpaths
    processBPs :: Value -> [BufferPath] -> (Value, [ByteString], [BufferPath])
processBPs Value
val = forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl (Value, [ByteString], [BufferPath])
-> BufferPath -> (Value, [ByteString], [BufferPath])
f (Value
val,[],[])
      where
#if MIN_VERSION_aeson(2,0,0)
        nestedLookupRemove :: BufferPath -> Value -> (Value, Maybe Value)
        nestedLookupRemove :: BufferPath -> Value -> (Value, Maybe Value)
nestedLookupRemove [] Value
v = (Value
v,forall a. a -> Maybe a
Just Value
v)
        nestedLookupRemove [Text
b] Value
v =
          case Value
v of
            Object Object
o -> (Object -> Value
Object forall a b. (a -> b) -> a -> b
$ forall v. Key -> KeyMap v -> KeyMap v
KM.delete (Text -> Key
Key.fromText Text
b) Object
o, forall v. Key -> KeyMap v -> Maybe v
KM.lookup (Text -> Key
Key.fromText Text
b) Object
o)
            Value
_ -> (Value
v, forall a. Maybe a
Nothing)
        nestedLookupRemove (Text
b:BufferPath
bp) Value
v =
          case Value
v of
            Object Object
o -> forall b a. b -> (a -> b) -> Maybe a -> b
maybe (Value
v,forall a. Maybe a
Nothing) ((Value, Maybe Value) -> (Value, Maybe Value)
upd forall b c a. (b -> c) -> (a -> b) -> a -> c
. BufferPath -> Value -> (Value, Maybe Value)
nestedLookupRemove BufferPath
bp) (forall v. Key -> KeyMap v -> Maybe v
KM.lookup (Text -> Key
Key.fromText Text
b) Object
o)
            Value
_ -> (Value
v,forall a. Maybe a
Nothing)
            where upd :: (Value, Maybe Value) -> (Value, Maybe Value)
                  upd :: (Value, Maybe Value) -> (Value, Maybe Value)
upd (Object Object
v', Just (Object Object
u)) = (Object -> Value
Object forall a b. (a -> b) -> a -> b
$ forall v. Key -> v -> KeyMap v -> KeyMap v
KM.insert (Text -> Key
Key.fromText Text
b) (Object -> Value
Object Object
u) Object
v', forall a. a -> Maybe a
Just forall a b. (a -> b) -> a -> b
$ Object -> Value
Object Object
u)
                  upd (Value, Maybe Value)
r = (Value, Maybe Value)
r
#else
        nestedLookupRemove :: BufferPath -> Value -> (Value, Maybe Value)
        nestedLookupRemove [] v = (v,Just v)
        nestedLookupRemove [b] v =
          case v of
            Object o -> (Object $ HM.delete b o, HM.lookup b o)
            _ -> (v, Nothing)
        nestedLookupRemove (b:bp) v =
          case v of
            Object o -> maybe (v,Nothing) (upd . nestedLookupRemove bp) (HM.lookup b o)
            _ -> (v,Nothing)
            where upd :: (Value, Maybe Value) -> (Value, Maybe Value)
                  upd (Object v', Just (Object u)) = (Object $ HM.insert b (Object u) v', Just $ Object u)
                  upd r = r
#endif

        f :: (Value, [ByteString], [BufferPath]) -> BufferPath -> (Value, [ByteString], [BufferPath])
        f :: (Value, [ByteString], [BufferPath])
-> BufferPath -> (Value, [ByteString], [BufferPath])
f r :: (Value, [ByteString], [BufferPath])
r@(Value
v,[ByteString]
bs,[BufferPath]
bps) BufferPath
bp =
          case BufferPath -> Value -> (Value, Maybe Value)
nestedLookupRemove BufferPath
bp Value
v of
            (Value
newv, Just (String Text
b)) -> (Value
newv, ByteString -> ByteString
B64.decodeLenient (Text -> ByteString
encodeUtf8 Text
b) forall a. a -> [a] -> [a]
: [ByteString]
bs, BufferPath
bpforall a. a -> [a] -> [a]
:[BufferPath]
bps)
            (Value, Maybe Value)
_ -> (Value, [ByteString], [BufferPath])
r

-- Override toJSON for PublishDisplayData for sending Display messages through [method .= custom]
data WidgetDisplay = WidgetDisplay MessageHeader [DisplayData]

instance ToJSON WidgetDisplay where
  toJSON :: WidgetDisplay -> Value
toJSON (WidgetDisplay MessageHeader
replyHeader [DisplayData]
ddata) =
    let pbval :: Value
pbval = forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ MessageHeader -> [DisplayData] -> Maybe Transient -> Message
PublishDisplayData MessageHeader
replyHeader [DisplayData]
ddata forall a. Maybe a
Nothing
    in forall a. ToJSON a => a -> Value
toJSON forall a b. (a -> b) -> a -> b
$ MessageHeader -> Value -> MessageType -> IPythonMessage
IPythonMessage MessageHeader
replyHeader Value
pbval MessageType
DisplayDataMessage

data IPythonMessage = IPythonMessage MessageHeader Value MessageType

instance ToJSON IPythonMessage where
  toJSON :: IPythonMessage -> Value
toJSON (IPythonMessage MessageHeader
replyHeader Value
val MessageType
mtype) =
    [Pair] -> Value
object
      [ Key
"header" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= MessageHeader
replyHeader
      , Key
"parent_header" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= String -> String
str String
""
      , Key
"metadata" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= [Pair] -> Value
object []
      , Key
"content" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= Value
val
      , Key
"msg_type" forall kv v. (KeyValue kv, ToJSON v) => Key -> v -> kv
.= (forall a. ToJSON a => a -> Value
toJSON forall b c a. (b -> c) -> (a -> b) -> a -> c
. MessageType -> String
showMessageType forall a b. (a -> b) -> a -> b
$ MessageType
mtype)
      ]

str :: String -> String
str :: String -> String
str = forall a. a -> a
id

-- Handle messages one-by-one, while updating state simultaneously
widgetHandler :: (Message -> IO ())
              -> MessageHeader
              -> KernelState
              -> [WidgetMsg]
              -> IO KernelState
widgetHandler :: (Message -> IO ())
-> MessageHeader -> KernelState -> [WidgetMsg] -> IO KernelState
widgetHandler Message -> IO ()
sender MessageHeader
hdr = forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldM ((Message -> IO ())
-> MessageHeader -> KernelState -> WidgetMsg -> IO KernelState
handleMessage Message -> IO ()
sender MessageHeader
hdr)