{-# LANGUAGE OverloadedStrings #-}
module Ghci.Websockets(
  -- $docs
  --
    initialise
  , initialiseDef
  , broadcast
  -- * Configuration
  , Config(..)
  , Verbosity(..)
  , defaultConfig
  ) where

import           Control.Concurrent (MVar, forkIO, modifyMVar, modifyMVar_,
                                     newMVar, readMVar)
import           Control.Exception  (catch)
import           Control.Monad      (void, (>=>))
import           Data.Aeson         (ToJSON)
import qualified Data.Aeson         as Aeson
import           Data.Foldable      (traverse_)
import qualified Data.Map           as Map
import qualified Data.Text          as Text
import qualified Foreign.Store      as Store
import qualified Network.WebSockets as WS

-- $docs
-- This modules implements a websocket server whose state survives GHCi 
-- reloads. To use it, run 'initialiseDef' once  per GHCi session, and then 
-- call 'broadcast' to send a JSON value to all clients that are currently 
-- connected. All messages from clients are ignored.

-- | What to do with log messages
data Verbosity =
      Verbose -- ^ Write all log messages to stdout
    | Silent -- ^ Ignore all log messages
    deriving (Eq, Ord, Show)

-- | Server configuration
data Config =
  Config
    { port      :: Int -- ^ What port to start the server on
    , verbosity :: Verbosity -- ^ What to do with log messages
    }

logStr :: Config -> String -> IO ()
logStr c s = case verbosity c of
  Silent  -> pure ()
  Verbose -> putStrLn s

-- | Default config, use port 9160 and ignore all log messages.
defaultConfig :: Config
defaultConfig = Config 9160 Silent

newtype ServerState = ServerState { unServerState :: Map.Map ConnectionID WS.Connection }

type ConnectionID = Int

serverState :: ServerState
serverState = ServerState Map.empty

addConnection :: WS.Connection -> ServerState -> (ServerState, ConnectionID)
addConnection c (ServerState mp) = (ServerState (Map.insert k c mp), k) where
  k = maybe 0 (succ . fst) $ Map.lookupMax mp

deleteConnection :: ConnectionID -> ServerState -> ServerState
deleteConnection i (ServerState mp) = ServerState (Map.delete i mp)

theStore :: Store.Store (MVar ServerState)
theStore = Store.Store 0

-- | Send a JSON object to all clients. Throws an exception if 'initialise' has
--   not been run first.
broadcast :: ToJSON a => a -> IO ()
broadcast t = Store.withStore theStore (readMVar >=> go) where
  go s = traverse_ (`WS.sendTextData` msg) (unServerState s)
  msg = Aeson.encode t

-- | Start the websocket server using the default config (port 9160). Call once
--   per GHCi session.
initialiseDef :: IO ()
initialiseDef = initialise defaultConfig

-- | Start the websocket server using the port specified in the config. Call
--   once per GHCi session.
initialise :: Config -> IO ()
initialise c@Config{port=p} = do
  state <- newMVar serverState
  Store.writeStore theStore state
  logStr c ("Starting websocket server on port " ++ show p)
  void $ forkIO (WS.runServer "127.0.0.1" p (application c))

application :: Config -> WS.ServerApp
application conf pending = do
  state <- Store.readStore theStore
  conn <- WS.acceptRequest pending
  WS.forkPingThread conn 30
  connID <- modifyMVar state (pure . addConnection conn)
  logStr conf $ "Accepted connection " ++ show connID
  let go = (WS.receiveData conn >>= logStr conf . Text.unpack) >> go
  catch go (closeConnection conf connID)

closeConnection :: Config -> ConnectionID -> WS.ConnectionException -> IO ()
closeConnection conf connID ex = do
  logStr conf $ "Closing connection " ++ show connID ++ " due to " ++ show ex
  state <- Store.readStore theStore
  modifyMVar_ state (pure . deleteConnection connID)