{-# LANGUAGE GADTs             #-}
{-# LANGUAGE OverloadedStrings #-}
-- |
-- Copyright : (c) NASA, 2024-2025
-- License   : BSD-style (see the LICENSE file in the distribution)
--
-- Run a Copilot simulation live and allow interacting with it using a
-- websocket.
--
-- This visualizer enables adding new streams to a visualization. To do so, the
-- visualizer needs access to the original spec, and needs to be able to
-- interpret new expressions in the same context as the prior expressions.
--
-- An example of a spec that can be passed as argument to the visualizer
-- follows:
--
-- @
--   spec :: String
--   spec = unlines
--     [ "let temperature :: Stream Word8"
--     , "    temperature = extern \"temperature\" (Just [0, 15, 20, 25, 30])"
--     , ""
--     , "    ctemp :: Stream Float"
--     , "    ctemp = (unsafeCast temperature) * (150.0 / 255.0) - 50.0"
--     , ""
--     , "    trueFalse :: Stream Bool"
--     , "    trueFalse = [True] ++ not trueFalse"
--     , ""
--     , "in do trigger \"heaton\"  (temperature < 18) [arg ctemp, arg (constI16 1), arg trueFalse]"
--     , "      trigger \"heatoff\" (temperature > 21) [arg (constI16 1), arg ctemp]"
--     , "      observer \"temperature\" temperature"
--     , "      observer \"temperature2\" (temperature + 1)"
--     ]
-- @
--
-- The imports are predefined.
module Copilot.Visualize.Live
    ( visualize
    , visualizeWith
    , VisualSettings(..)
    , mkDefaultVisualSettings
    , SimulationSettings(..)
    , mkDefaultSimulationSettings
    )
  where

-- External imports
import           Control.Exception  (SomeException (..), handle)
import           Data.Aeson         (ToJSON (..), encode)
import qualified Data.Text          as T
import qualified Network.WebSockets as WS
import           Prelude            hiding (div, not, (++), (<), (>))
import qualified Prelude
import           Text.Read          (readMaybe)

-- External imports: Copilot
import qualified Copilot.Core                   as Core
import           Copilot.Interpret.Eval         (ShowType (Haskell), eval)
import           Copilot.Language               hiding (interpret, typeOf)
import           Copilot.Visualize.UntypedTrace (AppData, makeTraceEval)

-- Internal imports
import Copilot.Visualize.Dynamic

-- | Start a simulation for an input spec, listening for commands and
-- communicating the status via a websocket.
visualize :: String -> IO ()
visualize :: String -> IO ()
visualize = VisualSettings -> String -> IO ()
visualizeWith VisualSettings
mkDefaultVisualSettings

-- | Start a simulation for an input spec, listening for commands and
-- communicating the status via a websocket.
visualizeWith :: VisualSettings -> String -> IO ()
visualizeWith :: VisualSettings -> String -> IO ()
visualizeWith VisualSettings
settings String
spec = do
  String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ [String] -> String
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat
    [ String
"WebSocket server starting on port "
    , Int -> String
forall a. Show a => a -> String
show (VisualSettings -> Int
visualSettingsPort VisualSettings
settings)
    , String
"..."
    ]
  String -> Int -> ServerApp -> IO ()
WS.runServer
    (VisualSettings -> String
visualSettingsHost VisualSettings
settings)
    (VisualSettings -> Int
visualSettingsPort VisualSettings
settings)
    (VisualSettings -> String -> ServerApp
app VisualSettings
settings String
spec)

-- | Settings used to customize the simulation and interaction.
data VisualSettings = VisualSettings
  { VisualSettings -> String
visualSettingsHost       :: String
                                -- ^ Host interface to listen to. Use
                                -- "127.0.0.1" to listen at localhost.

  , VisualSettings -> Int
visualSettingsPort       :: Int
                                -- ^ Port to listen to.

  , VisualSettings -> SimulationSettings
visualSettingsSimulation :: SimulationSettings
                                -- ^ Settings for the simulation.
  }

-- | Default settings that simulates 3 steps and listens on localhost at port
-- 9160.
mkDefaultVisualSettings :: VisualSettings
mkDefaultVisualSettings :: VisualSettings
mkDefaultVisualSettings = VisualSettings
  { visualSettingsHost :: String
visualSettingsHost       = String
"127.0.0.1"
  , visualSettingsPort :: Int
visualSettingsPort       = Int
9160
  , visualSettingsSimulation :: SimulationSettings
visualSettingsSimulation = SimulationSettings
mkDefaultSimulationSettings
  }

-- * Server

-- | Server application using web sockets.
app :: VisualSettings -> String -> WS.ServerApp
app :: VisualSettings -> String -> ServerApp
app VisualSettings
settings String
spec PendingConnection
pending = do
  Connection
conn <- PendingConnection -> IO Connection
WS.acceptRequest PendingConnection
pending
  Connection -> Int -> IO () -> IO () -> IO ()
forall a. Connection -> Int -> IO () -> IO a -> IO a
WS.withPingThread Connection
conn Int
30 (() -> IO ()
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return ()) (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ VisualSettings -> String -> Connection -> IO ()
appInit VisualSettings
settings String
spec Connection
conn

-- | Initialize the backend.
appInit :: VisualSettings -> String -> WS.Connection -> IO ()
appInit :: VisualSettings -> String -> Connection -> IO ()
appInit VisualSettings
settings String
spec Connection
conn = (SomeException -> IO ()) -> IO () -> IO ()
forall e a. Exception e => (e -> IO a) -> IO a -> IO a
handle SomeException -> IO ()
appException (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    -- Initialize the simulation.
    SimData
simData  <- SimulationSettings -> String -> IO SimData
simInit (VisualSettings -> SimulationSettings
visualSettingsSimulation VisualSettings
settings) String
spec

    -- Communicate the current values of the trace, in JSON, via the web
    -- socket.
    let appData :: AppData
appData = Int -> Spec -> AppData
mkAppData (SimData -> Int
simSteps SimData
simData) (SimData -> Spec
simSpec SimData
simData)
    let samples :: ByteString
samples = Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$ AppData -> Value
forall a. ToJSON a => a -> Value
toJSON AppData
appData
    Connection -> ByteString -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendTextData Connection
conn ByteString
samples

    -- Start the application loop.
    VisualSettings -> Connection -> SimData -> IO ()
appMainLoop VisualSettings
settings Connection
conn SimData
simData

  where

    appException :: SomeException -> IO ()
    appException :: SomeException -> IO ()
appException SomeException
e = do
      String -> IO ()
putStrLn (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ String
"Error:" String -> String -> String
forall a. [a] -> [a] -> [a]
Prelude.++ SomeException -> String
forall a. Show a => a -> String
show SomeException
e
      String -> IO ()
forall a. HasCallStack => String -> a
error (String -> IO ()) -> String -> IO ()
forall a b. (a -> b) -> a -> b
$ SomeException -> String
forall a. Show a => a -> String
show SomeException
e

-- | Run the main application, repeatedly reading commands from a web socket
-- and returning results via the same web socket.
appMainLoop :: VisualSettings
            -> WS.Connection
            -> SimData
            -> IO ()
appMainLoop :: VisualSettings -> Connection -> SimData -> IO ()
appMainLoop VisualSettings
settings Connection
conn SimData
simData = do
  -- Read a command from the web socket.
  Maybe (Command, String)
cmdM <- String -> Maybe (Command, String)
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe (Command, String))
-> (Text -> String) -> Text -> Maybe (Command, String)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> String
T.unpack (Text -> Maybe (Command, String))
-> IO Text -> IO (Maybe (Command, String))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Connection -> IO Text
forall a. WebSocketsData a => Connection -> IO a
WS.receiveData Connection
conn

  -- Run a simulation step, if a command has been received.
  let simulationSettings :: SimulationSettings
simulationSettings = VisualSettings -> SimulationSettings
visualSettingsSimulation VisualSettings
settings
  SimData
simData' <- IO SimData
-> ((Command, String) -> IO SimData)
-> Maybe (Command, String)
-> IO SimData
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (SimData -> IO SimData
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure SimData
simData) (SimulationSettings -> SimData -> (Command, String) -> IO SimData
simStep SimulationSettings
simulationSettings SimData
simData) Maybe (Command, String)
cmdM

  -- Communicate the current values of the trace, in JSON, via the web socket.
  let appData :: AppData
appData = Int -> Spec -> AppData
mkAppData (SimData -> Int
simSteps SimData
simData') (SimData -> Spec
simSpec SimData
simData')
      samples :: ByteString
samples = Value -> ByteString
forall a. ToJSON a => a -> ByteString
encode (Value -> ByteString) -> Value -> ByteString
forall a b. (a -> b) -> a -> b
$ AppData -> Value
forall a. ToJSON a => a -> Value
toJSON AppData
appData
  Connection -> ByteString -> IO ()
forall a. WebSocketsData a => Connection -> a -> IO ()
WS.sendTextData Connection
conn ByteString
samples

  VisualSettings -> Connection -> SimData -> IO ()
appMainLoop VisualSettings
settings Connection
conn SimData
simData'

-- * Auxiliary functions

-- | Obtain the trace data from a Copilot spec for a number of steps.
mkAppData :: Int -> Core.Spec -> AppData
mkAppData :: Int -> Spec -> AppData
mkAppData Int
numSteps Spec
spec' =
  Int -> Spec -> ExecTrace -> AppData
makeTraceEval Int
numSteps Spec
spec' (ShowType -> Int -> Spec -> ExecTrace
eval ShowType
Haskell Int
numSteps Spec
spec')