content: {-
     Adapted from haskell-chat-sever-example which is
        Copyright (c) 2012, Joseph Adams
  
     Modifications (c) 2012, Simon Marlow
  -}
  
  {-# LANGUAGE RecordWildCards #-}
  module Main where
  
  import ConcurrentUtils
  
  import Control.Concurrent
  import Control.Concurrent.STM
  import Control.Concurrent.Async
  import qualified Data.Map as Map
  import Data.Map (Map)
  import System.IO
  import Control.Exception
  import NetworkUtils
  import Control.Monad
  import Text.Printf
  
  {-
  Notes
  
  - protocol:
      Server: "Name?"
      Client: <string>
      -- if <string> is already in use, ask for another name
      -- Commands:
      --   /tell <string> message...  (single-user tell)
      --   /quit                      (exit)
      --   /kick <string>             (kick another user)
      --   message...                 (broadcast to all connected users)
  
  - a client needs to both listen for commands from the socket and
    listen for activity from other clients.  Therefore we're going to
    need at least two threads per client (for listening to multiple
    things).  Easiest is to use STM for in-process communication, and to
    have a receiving thread that listens on the socket and forwards to a
    TChan.
  
  - Handle all errors properly, be async-exception safe
  
  - Consistency:
    - if two clients simultaneously kick a third client, only one will be
      successful
  
  See doc/lab-exercises.tex for some ideas for enhancements that you
  could try.
  
  -}
  
  -- <<main
  main :: IO ()
  main = withSocketsDo $ do
    server <- newServer
    listenOn port $ \sock -> do
      printf "Listening on port %d\n" port
      forever $
        accept sock $ \(handle, peer) -> do
          printf "Accepted connection from %s\n" (show peer)
          forkFinally (talk handle server) (\_ -> hClose handle)
  
  port :: Int
  port = 44444
  -- >>
  
  
  -- ---------------------------------------------------------------------------
  -- Data structures and initialisation
  
  -- <<Client
  type ClientName = String
  
  data Client = Client
    { clientName     :: ClientName
    , clientHandle   :: Handle
    , clientKicked   :: TVar (Maybe String)
    , clientSendChan :: TChan Message
    }
  -- >>
  
  -- <<newClient
  newClient :: ClientName -> Handle -> STM Client
  newClient name handle = do
    c <- newTChan
    k <- newTVar Nothing
    return Client { clientName     = name
                  , clientHandle   = handle
                  , clientSendChan = c
                  , clientKicked   = k
                  }
  -- >>
  
  -- <<Server
  data Server = Server
    { clients :: TVar (Map ClientName Client)
    }
  
  newServer :: IO Server
  newServer = do
    c <- newTVarIO Map.empty
    return Server { clients = c }
  -- >>
  
  -- <<Message
  data Message = Notice String
               | Tell ClientName String
               | Broadcast ClientName String
               | Command String
  -- >>
  
  -- -----------------------------------------------------------------------------
  -- Basic operations
  
  -- <<broadcast
  broadcast :: Server -> Message -> STM ()
  broadcast Server{..} msg = do
    clientmap <- readTVar clients
    mapM_ (\client -> sendMessage client msg) (Map.elems clientmap)
  -- >>
  
  -- <<sendMessage
  sendMessage :: Client -> Message -> STM ()
  sendMessage Client{..} msg =
    writeTChan clientSendChan msg
  -- >>
  
  -- <<sendToName
  sendToName :: Server -> ClientName -> Message -> STM Bool
  sendToName server@Server{..} name msg = do
    clientmap <- readTVar clients
    case Map.lookup name clientmap of
      Nothing     -> return False
      Just client -> sendMessage client msg >> return True
  -- >>
  
  tell :: Server -> Client -> ClientName -> String -> IO ()
  tell server@Server{..} Client{..} who msg = do
    ok <- atomically $ sendToName server who (Tell clientName msg)
    if ok
       then return ()
       else hPutStrLn clientHandle (who ++ " is not connected.")
  
  kick :: Server -> ClientName -> ClientName -> STM ()
  kick server@Server{..} who by = do
    clientmap <- readTVar clients
    case Map.lookup who clientmap of
      Nothing ->
        void $ sendToName server by (Notice $ who ++ " is not connected")
      Just victim -> do
        writeTVar (clientKicked victim) $ Just ("by " ++ by)
        void $ sendToName server by (Notice $ "you kicked " ++ who)
  
  -- -----------------------------------------------------------------------------
  -- The main server
  
  talk :: Handle -> Server -> IO ()
  talk handle server@Server{..} = do
    hSetNewlineMode handle universalNewlineMode
        -- Swallow carriage returns sent by telnet clients
    hSetBuffering handle LineBuffering
    readName
   where
  -- <<readName
    readName = do
      hPutStrLn handle "What is your name?"
      name <- hGetLine handle
      if null name
        then readName
        else mask $ \restore -> do        -- <1>
               ok <- checkAddClient server name handle
               case ok of
                 Nothing -> restore $ do  -- <2>
                    hPrintf handle
                       "The name %s is in use, please choose another\n" name
                    readName
                 Just client ->
                    restore (runClient server client) -- <3>
                        `finally` removeClient server name
  -- >>
  
  -- <<checkAddClient
  checkAddClient :: Server -> ClientName -> Handle -> IO (Maybe Client)
  checkAddClient server@Server{..} name handle = atomically $ do
    clientmap <- readTVar clients
    if Map.member name clientmap
      then return Nothing
      else do client <- newClient name handle
              writeTVar clients $ Map.insert name client clientmap
              broadcast server  $ Notice (name ++ " has connected")
              return (Just client)
  -- >>
  
  -- <<removeClient
  removeClient :: Server -> ClientName -> IO ()
  removeClient server@Server{..} name = atomically $ do
    modifyTVar' clients $ Map.delete name
    broadcast server $ Notice (name ++ " has disconnected")
  -- >>
  
  -- <<runClient
  runClient :: Server -> Client -> IO ()
  runClient serv@Server{..} client@Client{..} = do
    race server receive
    return ()
   where
    receive = forever $ do
      msg <- hGetLine clientHandle
      atomically $ sendMessage client (Command msg)
  
    server = join $ atomically $ do
      k <- readTVar clientKicked
      case k of
        Just reason -> return $
          hPutStrLn clientHandle $ "You have been kicked: " ++ reason
        Nothing -> do
          msg <- readTChan clientSendChan
          io <- handleMessage serv client msg
          return $ do
              continue <- io
              when continue $ server
  -- >>
  
  -- <<handleMessage
  handleMessage :: Server -> Client -> Message -> STM (IO Bool)
  handleMessage server client@Client{..} message =
    case message of
       Notice msg         -> output $ "*** " ++ msg
       Tell name msg      -> output $ "*" ++ name ++ "*: " ++ msg
       Broadcast name msg -> output $ "<" ++ name ++ ">: " ++ msg
       Command msg ->
         case words msg of
             ["/kick", who] -> do
                 kick server who clientName
                 return $ return True
             "/tell" : who : what -> do
                 return $ do
                   tell server client who (unwords what)
                   return True
             ["/quit"] ->
                 return $ return False
             ('/':_):_ -> return $ do
                 hPutStrLn clientHandle $ "Unrecognised command: " ++ msg
                 return True
             _ -> do
                 broadcast server $ Broadcast clientName msg
                 return $ return True
   where
     output s = return $ do hPutStrLn clientHandle s; return True
  -- >>
  