module Main where

import Network

import IO

import System.IO

import Control.Concurrent

import Control.Concurrent.STM

import Control.Exception (finally)

import Control.Monad (forM, filterM, liftM)

import Miscutils

type Client = (TChan String, Handle)

acceptLoop :: Socket -> TChan Client -> IO ()

acceptLoop servSock acceptChan = do

	(cHandle, host, port) <- accept servSock

	cChan <- atomically newTChan

	forkIO $ clientLoop cHandle cChan

	atomically $ writeTChan acceptChan (cChan, cHandle)

	acceptLoop servSock acceptChan

listenLoop :: Handle -> TChan String -> IO ()

listenLoop handle chan = do

	str <- hGetLine handle

	atomically $ writeTChan chan str

	listenLoop handle chan

clientLoop :: Handle -> TChan String -> IO ()

clientLoop handle chan =

	listenLoop handle chan

		`catch` (const $ clientOff >> return ())

		`finally` hClose handle

	where clientOff = atomically $ writeTChan chan "QUIT"

mainLoop :: Socket -> TChan Client -> [Client] -> IO ()

mainLoop servSock acceptChan clients = do

	r <- atomically $ (Left `fmap` readTChan acceptChan) `orElse` (Right `fmap` tselect clients)

	case r of

		Left (ch, h) -> do

			mainLoop servSock acceptChan $ (ch, h):clients

		Right (line, handle) -> do

			clients' <- forM clients $

					\(ch, h) -> do

						hPutStrLn h line

						hFlush h

						return [(ch,h)]

					`catch` const (hClose h >> return [])

			mainLoop servSock acceptChan $ concat clients'

tselect :: [(TChan a, t)] -> STM (a, t)

tselect = foldl orElse retry . map (\(ch, ty) -> (flip (,) ty) `fmap` readTChan ch)

startServer serverSocket = do

	acceptChan <- atomically newTChan

	forkIO $ acceptLoop serverSocket acceptChan

	mainLoop serverSocket acceptChan []

main = withSocketsDo $ do

	serverSocket <- listenOn $ Service "hedgewars"

	startServer serverSocket `finally` sClose serverSocket