{-# LANGUAGE DerivingVia #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{-# OPTIONS_GHC -Wno-unrecognised-pragmas #-}

{-# HLINT ignore "Redundant bracket" #-}

module WebSocket (broadCastToClients, 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)

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