{-# LANGUAGE DerivingVia          #-}
{-# LANGUAGE FlexibleContexts     #-}
{-# LANGUAGE UndecidableInstances #-}


module WebSocket (broadCastToClients, runWebSocketServer) where

import           ClassyPrelude
import           Data.UUID             (UUID)
import           Data.UUID.V4          (nextRandom)
import qualified Network.WebSockets    as WS
import           Types.AppTypes        (Env (connectedClientsState),
                                        HasConnectedClientState (getConnectedClientState))
import           Types.ConnectionState (Client (..), ConnectedClients)

data WSEnv = WSEnv
  { appEnv     :: Env,
    connection :: WS.Connection,
    clientId   :: UUID
  }

instance HasConnectedClientState WSEnv where
  getConnectedClientState = connectedClientsState . appEnv

class HasWSEnv a where
  getConn :: a -> WS.Connection
  getClientId :: a -> UUID

instance HasWSEnv WSEnv where
  getConn = connection
  getClientId = clientId

addWSClient ::
  ( HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  Client ->
  m ()
addWSClient client = do
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s -> do
    let s' = addClient client s
    return s'

addClient :: Client -> ConnectedClients -> ConnectedClients
addClient client clients = client : clients

disconnectWsClient ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
disconnectWsClient = do
  clientId <- getClientId <$> ask
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s ->
    let s' = removeClient clientId s
     in return s'

joinRoom ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
joinRoom = do
  clientId <- getClientId <$> ask
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s ->
    let s' = updateClient clientId (\c -> c {joinedRoom = True}) s
     in return s'

removeClient :: UUID -> ConnectedClients -> ConnectedClients
removeClient toRemove = filter ((/= toRemove) . uuid)

updateClient :: UUID -> (Client -> Client) -> ConnectedClients -> ConnectedClients
updateClient toUpdate f allClients = do
  client <- allClients
  if (toUpdate == uuid client)
    then (return $ f client)
    else return client

broadcast :: Text -> ConnectedClients -> IO ()
broadcast message clients = do
  putStrLn message
  forM_ clients $ \client -> WS.sendTextData (conn client) message

runWebSocketServer ::
  ( MonadIO m,
    MonadReader Env m
  ) =>
  m ()
runWebSocketServer = do
  putStrLn "Websocket up at 127.0.0.1:9160"
  wsApp' <- runWSApp
  liftIO $ WS.runServer "127.0.0.1" 9160 wsApp'

runWSApp ::
  ( MonadIO m,
    MonadReader Env m
  ) =>
  m WS.ServerApp
runWSApp = do
  env <- ask
  return
    ( \pending -> do
        conn <- WS.acceptRequest pending
        uuid <- nextRandom
        let wsEnv = WSEnv {appEnv = env, connection = conn, clientId = uuid}
        WS.withPingThread conn 30 (return ()) $ do
          unWSApp wsApp wsEnv
    )

newClient ::
  ( MonadIO m,
    HasWSEnv env,
    MonadReader env m
  ) =>
  Text ->
  m Client
newClient name = do
  env <- ask
  return $ Client {uuid = getClientId env, name = name, conn = getConn env, joinedRoom = False}

newtype WSApp env a = WSApp {unWSApp :: env -> IO a}
  deriving
    ( Functor,
      Applicative,
      Monad,
      MonadReader env,
      MonadIO,
      MonadUnliftIO
    )
    via ReaderT env IO

wsApp ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
wsApp = do
  msg <- getMessage
  putStrLn msg
  client <- newClient msg
  addWSClient client
  flip finally disconnectWsClient $ do
    forever $ do
      currentMsg <- getMessage
      joinRoom
      putStrLn currentMsg

getMessage ::
  ( HasWSEnv env,
    MonadIO m,
    MonadReader env m
  ) =>
  m Text
getMessage = do
  conn' <- getConn <$> ask
  liftIO $ WS.receiveData conn'

broadCastToClients ::
  ( MonadIO m,
    HasConnectedClientState env,
    MonadReader env m
  ) =>
  Text ->
  m ()
broadCastToClients message = do
  state <- getConnectedClientState <$> ask
  liftIO $ withMVar state $ \currenState -> broadcast message currenState