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

module WebSocket (broadcast, initMVarState, runWebSocketServer) where

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

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

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

newConnectedClients :: ConnectedClients
newConnectedClients = []

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

initMVarState :: IO (MVar ConnectedClients)
initMVarState = newMVar newConnectedClients

runWebSocketServer ::
  ( MonadIO m,
    MonadReader Env m
  ) =>
  m ()
runWebSocketServer = do
  putStrLn "Websocket up at 127.0.0.1:9160"
  state <- getConnectedClientState <$> ask
  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
        putStrLn "pending request"
        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}

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

-- instance MonadBaseControl IO m => MonadBaseControl IO (WSApp env n)
-- where

wsApp ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  m ()
wsApp = do
  state <- getConnectedClientState <$> ask
  msg <- getMessage
  putStrLn msg
  client <- newClient msg
  modifyMVar_ state $ \s -> do
    let s' = addClient client s
    return s'
  flip finally disconnectWsClient $ do
    forever $ do
      currentMsg <- getMessage
      putStrLn currentMsg

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

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

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