import System.Environment ( getArgs )
import Network.Socket
import System.IO ( putStrLn )
import System.IO.Error ( isEOFError )
import Control.Concurrent ( forkIO )
import Control.Monad ( forever, forM_, when )
import Control.Monad.Trans ( liftIO )
import Control.Monad.State ( StateT, evalStateT, get )
import Control.Concurrent.STM
    ( TChan
    , newTChan
    , writeTChan
    , readTChan
    , atomically
    , TVar
    , newTVar
    , readTVar
    , writeTVar
    )
import qualified Data.Map as Map

import Network.Chat.Protocol
import Network.Chat.Parse ( parseClientMessage )
import Network.Chat.Util ( while )

-- |Represents a network endpoint.
data Endpoint = Endpoint { epSocket :: Socket
                         , epConnected :: Bool
                         }

-- |Represents an initialized chat client.
data Client = Client { clEndpoint :: Endpoint
                     , clUsername :: Username
                     , clChannel :: TChan ServerMessage
                     }

-- |The type of transactional user-to-client mapping that the server
-- |uses to route messages between clients.
type Clients = TVar (Map.Map Username Client)

-- |The state environment in which most server code runs.
type ServerEnv a = StateT Clients IO a

-- |Given a client username, get the Client from the environment or
-- |Nothing if no such client is connected.
getClient :: Username -> ServerEnv (Maybe Client)
getClient u = do
  clients <- get
  clientMap <- liftIO $ atomically $ readTVar clients
  return $ Map.lookup u clientMap

-- |Given a Client, add it to the environment.
addClient :: Client -> ServerEnv ()
addClient client = do
  clients <- get
  clientMap <- liftIO $ atomically $ readTVar clients
  liftIO $ atomically $ writeTVar clients $ Map.insert (clUsername client) client clientMap
  liftIO $ putStrLn $ "Added client: " ++ (show (clUsername client))

-- |Given a client username, if there is a client in the environment,
-- remove it.
removeClient :: Username -> ServerEnv ()
removeClient u = do
  clients <- get
  clientMap <- liftIO $ atomically $ readTVar clients
  liftIO $ atomically $ writeTVar clients $ Map.delete u clientMap
  liftIO $ putStrLn $ "Removed client: " ++ (show u)

-- |Is the client endpoint reachable over the network?  True if the
-- client has not sent a Disconnect message.
isConnected :: Username -> ServerEnv Bool
isConnected u = do
  client <- getClient u
  case client of
    Nothing -> return False
    Just c -> return $ epConnected (clEndpoint c)

-- |Disconnect a network endpoint.
disconnectEp :: Endpoint -> IO ()
disconnectEp endpoint = do
  sClose (epSocket endpoint)

-- |Disconnect a client by disconnecting its network endpoint.  Remove
-- the client from the environment.
disconnectClient :: Username -> ServerEnv ()
disconnectClient u = do
  client <- getClient u
  case client of
    Nothing -> return ()
    Just c -> do
      let ep = clEndpoint c
      liftIO $ disconnectEp ep
      removeClient u

-- |Given a server message and client, serialize the message send it
-- to the client over the client's network endpoint.
clientSendMessage :: ServerMessage -> Client -> IO Int
clientSendMessage msg client = do
  send (epSocket (clEndpoint client)) (show msg)

-- |Serialize a protocol message and send it to the client with the
-- specified username.  Returns the number of bytes sent.
sendMessage :: ServerMessage -> Username -> ServerEnv Int
sendMessage msg destUser = do
  client <- getClient destUser
  case client of
    Nothing -> return 0
    Just c -> liftIO $ clientSendMessage msg c

-- |Given a server message, sent it to every connected client.
broadcast :: ServerMessage -> ServerEnv ()
broadcast msg = do
  clients <- get
  clientMap <- liftIO $ atomically $ readTVar clients
  forM_ (Map.keys clientMap) (sendMessage msg)

-- |Given a client username, monitor the specified client's message
-- channel as long as the client is connected; when a message arrives
-- in the channel, sent it to the client over its network endpoint.
-- The channel will receive messages from other clients.
clientChanMonitor :: Username -> ServerEnv ()
clientChanMonitor u = do
  -- Read a message from the client's output channel, and serialize it
  -- and send it to the client over the network.
  while (isConnected u) $ do
    client <- getClient u
    case client of
      Nothing -> do
                  liftIO $ putStrLn $ "clientChanMonitor terminating for " ++ (show u)
                  return ()
      Just c -> do
        msg <- liftIO $ atomically $ readTChan (clChannel c)
        sendMessage msg u
        return ()

-- |Listen for a message from the client; if the message can be
-- parsed, return a corresponding ClientMessage or Nothing if it
-- cannot be parsed.
nextMessage :: Endpoint -> IO ClientMessage
nextMessage ep = do
  message <- recv (epSocket ep) 1024
  case parseClientMessage message of
    Nothing -> nextMessage ep
    Just m -> return m

-- |Given a client username and a server message, queue the message on
-- the specified client's outgoing message queue.
queueMessage :: Username -> ServerMessage -> ServerEnv ()
queueMessage u msg = do
  client <- getClient u
  case client of
    Nothing -> return ()
    Just c -> liftIO $ atomically $ writeTChan (clChannel c) msg

-- |Given a protocol message from the specified user, perform the
-- appropriate action.
handleMessage :: Username -> ClientMessage -> ServerEnv ()
handleMessage u (Disconnect) = disconnectClient u >> broadcast (AnnounceDisconnect u)
handleMessage u (MessageTo destUser msg) = queueMessage destUser (MessageFrom u msg)
handleMessage u (Identify _) = do
  queueMessage u $ Error $ "Already identified as: " ++ (show u)
handleMessage u (GetUserList) = do
  clients <- get
  users <- liftIO $ atomically $ do
                                  clientMap <- readTVar clients
                                  return $ Map.keys clientMap
  queueMessage u $ UserList users

-- |Handle a new client connection by processing messages from that
-- client as long as it is connected.
handleClientConnection :: Username -> ServerEnv ()
handleClientConnection u = do
  while (isConnected u) $ do
    client <- getClient u
    case client of
      Nothing -> do
                  liftIO $ putStrLn $ "handleClientConnection terminating for " ++ (show u)
                  return ()
      Just c -> do
        m <- liftIO $ nextMessage (clEndpoint c)
        handleMessage u m

-- |Handler for EOF errors when using sockets.
handleClosedConnection :: IOError -> IO ()
handleClosedConnection e = do
  putStrLn $ "[connection closed] " ++ show e

-- |Registry of error handlers for various types of IOError.
errorHandlers :: [((IOError -> Bool), (IOError -> IO ()))]
errorHandlers = [ ( isEOFError, handleClosedConnection ) ]

-- |Handle IOErrors by looking at errorHandlers or, if an error has no
-- registered handler, just print an error message to the console.
errorHandler :: IOError -> IO ()
errorHandler e =
    let allHandlers = errorHandlers ++ [(\_ -> True, defaultHandler)]
        matchingHandlers = filter (\(matcher, _) -> matcher e) allHandlers
        defaultHandler err = putStrLn $ "[error] " ++ show err
    in do
      (snd (head matchingHandlers)) e

-- |Given a socket, run the server's main loop on the socket,
-- accepting new client connections and handling exceptions.
serveOn :: Socket -> Clients -> IO ()
serveOn listenSock serverEnv = forever $ do
  (clientSock, _) <- accept listenSock
  let ep = Endpoint clientSock True
  forkIO $ do
    client <- initClient ep serverEnv
    case client of
      Nothing -> return ()
      Just c -> do
             -- announce connection
             evalStateT (broadcast (AnnounceConnect (clUsername c))) serverEnv
             -- add the client to the server state
             evalStateT (addClient c) serverEnv
             -- begin accepting messages from the client
             startServingClient c serverEnv

-- |Begin serving the specified client in the specified server
-- environment.
startServingClient :: Client -> Clients -> IO ()
startServingClient client serverEnv = do
  -- now that the client has been initialized, forkIO a monitor for
  -- its communication channel.
  let u = clUsername client
  forkIO $ (evalStateT (clientChanMonitor u) serverEnv) `catch`
         (\e -> do
            evalStateT (broadcast (AnnounceDisconnect u)) serverEnv
            errorHandler e)
  -- Handle post-initialization messages from the client until it
  -- disconnects.
  evalStateT (handleClientConnection u) serverEnv `catch`
         (\e -> do
            evalStateT (broadcast (AnnounceDisconnect u)) serverEnv
            errorHandler e)

-- |Initialize a client by reading initialization messages from the
-- specified network endpoint.  Return a Client if initialization
-- succeeds, or Nothing if an error occurs or the endpoint is
-- disconnected.
initClient :: Endpoint -> Clients -> IO (Maybe Client)
initClient endpoint otherClients = do
  msg <- nextMessage endpoint
  case msg of
    -- Disconnect is a little special in the pre-initialization case
    -- since it means "give up."
    Disconnect -> disconnectEp endpoint >> return Nothing
    Identify u -> do
             channel <- atomically $ newTChan
             let client = Client endpoint u channel
             -- If a user with this name is already connected, send a
             -- Failure response and disconnect.
             connectedClients <- atomically $ readTVar otherClients
             case Map.lookup u connectedClients of
               Just _ -> do
                    clientSendMessage (IdentifyResult Failure u) client
                    return Nothing
               Nothing -> do
                    clientSendMessage (IdentifyResult Success u) client
                    return $ Just client
    -- No other messages make sense at this stage.  Maybe this means
    -- we should have two message types: initialization messages and
    -- post-initialization messages?
    _ -> initClient endpoint otherClients

-- |Given a hostname and port, create a listening socket bound to the
-- specified interface and port.
listenSocket :: String -> String -> IO Socket
listenSocket hostname port = do
    addrInfo <- getAddrInfo Nothing (Just hostname) (Just port)
    listenSock <- socket AF_INET Stream defaultProtocol
    let listenAddr = addrAddress (head addrInfo)
    bindSocket listenSock listenAddr
    listen listenSock 1
    return listenSock

main :: IO ()
main = do
  args <- getArgs
  when (length args /= 2) $ error "Usage: chat-server <hostname> <port>"
  let [hostname, port] = args

  withSocketsDo $ do
         listenSock <- listenSocket hostname port
         serverEnv <- atomically $ newTVar Map.empty
         serveOn listenSock serverEnv