{-# 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 = visualizeWith mkDefaultVisualSettings

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

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

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

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

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

-- * Server

-- | Server application using web sockets.
app :: VisualSettings -> String -> WS.ServerApp
app settings spec pending = do
  conn <- WS.acceptRequest pending
  WS.withPingThread conn 30 (return ()) $ appInit settings spec conn

-- | Initialize the backend.
appInit :: VisualSettings -> String -> WS.Connection -> IO ()
appInit settings spec conn = handle appException $ do
    -- Initialize the simulation.
    simData  <- simInit (visualSettingsSimulation settings) spec

    -- Communicate the current values of the trace, in JSON, via the web
    -- socket.
    let appData = mkAppData (simSteps simData) (simSpec simData)
    let samples = encode $ toJSON appData
    WS.sendTextData conn samples

    -- Start the application loop.
    appMainLoop settings conn simData

  where

    appException :: SomeException -> IO ()
    appException e = do
      putStrLn $ "Error:" Prelude.++ show e
      error $ show 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 settings conn simData = do
  -- Read a command from the web socket.
  cmdM <- readMaybe . T.unpack <$> WS.receiveData conn

  -- Run a simulation step, if a command has been received.
  let simulationSettings = visualSettingsSimulation settings
  simData' <- maybe (pure simData) (simStep simulationSettings simData) cmdM

  -- Communicate the current values of the trace, in JSON, via the web socket.
  let appData = mkAppData (simSteps simData') (simSpec simData')
      samples = encode $ toJSON appData
  WS.sendTextData conn samples

  appMainLoop settings conn simData'

-- * Auxiliary functions

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