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

module WebSocket (runWebSocketServer) where

import BroadcastUserData (broadcastUserData)
import ClassyPrelude
import Data.Aeson
import Data.UUID (UUID)
import Data.UUID.V4 (nextRandom)
import Network.WebSockets qualified as WS
import Types.AppTypes
  ( Env (..),
    HasConnectedClientState (getConnectedClientState),
  )
import Types.ConnectionState
  ( Client (..),
    ConnectedClients,
  )
import Types.RoomsState (HasRoomsState (..))
import Types.WebSocketMessages.WebSocketMessages
  ( SetClientInfo (displayName),
    WebSocketMessage (..),
  )

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

instance HasConnectedClientState WSEnv where
  getConnectedClientState = connectedClientsState . appEnv

instance HasRoomsState WSEnv where
  getRoomsState = roomsState . 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
  liftIO $ putStrLn "joinedRoom"
  modifyMVar_ state $ \s ->
    let s' = updateClient clientId (\c -> c {joinedRoom = True}) s
     in return s'

updateClientName ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    MonadReader env m,
    MonadUnliftIO m
  ) =>
  SetClientInfo ->
  m ()
updateClientName clientInfo = do
  clientId <- getClientId <$> ask
  state <- getConnectedClientState <$> ask
  modifyMVar_ state $ \s ->
    let s' = updateClient clientId (\c -> c {name = displayName clientInfo}) 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

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 ()) $ unWSApp wsApp wsEnv
    )

newClient ::
  ( MonadIO m,
    HasWSEnv env,
    MonadReader env m
  ) =>
  SetClientInfo ->
  m Client
newClient clientInfo = do
  env <- ask
  return $ Client {uuid = getClientId env, name = displayName clientInfo, 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

class Monad m => MonadWebSocket m where
  getTypedWSMessage :: FromJSON a => m a

instance MonadWebSocket (WSApp WSEnv) where
  getTypedWSMessage = do
    msg <- getMessage
    case eitherDecodeStrict $ encodeUtf8 msg of
      Right a -> return a
      Left err -> do
        sendMessage $ "Bad message: " <> pack err
        getTypedWSMessage

wsApp ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    HasRoomsState env,
    MonadReader env m,
    MonadUnliftIO m,
    MonadWebSocket m
  ) =>
  m ()
wsApp = do
  msg <- getTypedWSMessage
  print msg
  client <- newClient msg
  addWSClient client
  broadcastUserData
  flip finally disconnectWsClient $ forever $ do
    handleWSAction
    broadcastUserData

handleWSAction ::
  ( HasWSEnv env,
    HasConnectedClientState env,
    HasRoomsState env,
    MonadReader env m,
    MonadUnliftIO m,
    MonadWebSocket m
  ) =>
  m ()
handleWSAction = do
  msg <- getTypedWSMessage
  case msg of
    JoinRoomMessage _ -> do
      joinRoom
      return ()
    ClientInfoMessage clientInfo -> do
      updateClientName clientInfo

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

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