{-# LANGUAGE OverloadedStrings, TypeFamilies, ScopedTypeVariables #-}
module Network.GenericServer
  (genericServer,
   readUntilEmptyLine,
   waitData,
   Server (..),
   wrap,
   server
  ) where

import Control.Concurrent
import Control.Exception
import Control.Monad
import Network
import System.IO
import qualified Data.ByteString.Char8 as B

class Server s where
  type Request s
  type Response s

  toRequest :: s -> B.ByteString -> Request s
  fromResponse :: s -> Response s -> B.ByteString

  requestReader :: s -> Handle -> IO B.ByteString
  requestReader _ = readUntilEmptyLine

  worker :: s -> Request s -> IO (Response s, s)

  getServer :: s -> IO s
  getServer s = return s

  putServer :: s -> IO ()
  putServer _ = return ()

wrap :: (a -> IO b) -> (s -> a -> IO (b, s))
wrap fn = \s a -> do
    b <- fn a
    return (b, s)

server :: forall s. Server s => s -> Int -> IO ()
server s port = genericServer port (requestReader s) callOut
  where
    callOut str = do
       srv <- getServer s
       (res, srv') <- worker srv (toRequest srv str)
       putServer srv'
       return $ fromResponse srv' res

-- | Run TCP/IP server with any worker funciton
genericServer :: Int                                -- ^ Port number
              -> (Handle -> IO B.ByteString)        -- ^ Query reading function
              -> (B.ByteString -> IO B.ByteString)  -- ^ Worker function
              -> IO ()
genericServer port reader callOut = do
      sock  <- listenOn (PortNumber $ fromIntegral port)
      (forever $ loop sock) `finally` sClose sock
  where
    loop :: Socket -> IO ThreadId
    loop sock = do
      (h,_,_) <- accept sock
      forkIO $ do 
               hSetBuffering h NoBuffering
               text <- reader h
               result <- callOut text
               B.hPutStrLn h result
               hClose h

readUntilEmptyLine :: Handle -> IO B.ByteString
readUntilEmptyLine h = do
  str <- B.hGetLine h
  if (str == "\n") || (str == "\r") || (str == "\r\n")
    then return str
    else do
         next <- readUntilEmptyLine h
         return $ str `B.append` next

-- | Read given amount of bytes from socket
waitData :: Int             -- ^ Port number
         -> Int             -- ^ Data size
         -> IO B.ByteString
waitData port size = do
    sock  <- listenOn (PortNumber $ fromIntegral port)
    wait sock `finally` sClose sock
  where
    wait :: Socket -> IO B.ByteString
    wait s = do
      (h,_,_) <- accept s
      hSetBuffering h NoBuffering
      d <- B.hGet h size
      hClose h
      return d