module Main where

import Control.Concurrent (forkIO)
import Control.Concurrent.MVar (newEmptyMVar, putMVar, takeMVar)
import Control.Exception (bracket)
import Control.Monad (when)
import Network.Socket hiding (recv, recvFrom, send)
import System.Exit (exitFailure)
import Test.HUnit (Counts(..), Test(..), (@=?), runTestTT)

import qualified Data.ByteString as S
import qualified Data.ByteString.Char8 as C
import qualified Data.ByteString.Lazy.Char8 as L

import Network.Socket.ByteString (recv, recvFrom, send, sendAll, sendMany)
import qualified Network.Socket.ByteString.Lazy as NSBL

------------------------------------------------------------------------

port :: PortNumber
port = fromIntegral (3000 :: Int)

testMsg :: S.ByteString
testMsg = C.pack "This is a test message."

testLazySend :: Test
testLazySend = TestCase $ connectedTest client server
    where
      server sock = recv sock 1024 >>= (@=?) (C.take 1024 strictTestMsg)
      client sock = NSBL.send sock lazyTestMsg >>= (@=?) 1024

      -- message containing too many chunks to be sent in one system call
      lazyTestMsg = let alphabet = map C.singleton ['a'..'z']
                    in L.fromChunks (concat (replicate 100 alphabet))

      strictTestMsg = C.concat . L.toChunks $ lazyTestMsg

------------------------------------------------------------------------
-- Tests

testSendAll :: Test
testSendAll = TestCase $ connectedTest client server
    where
      server sock = recv sock 1024 >>= (@=?) testMsg
      client sock = sendAll sock testMsg

testSendMany :: Test
testSendMany = TestCase $ connectedTest client server
    where
      server sock = recv sock 1024 >>= (@=?) (S.append seg1 seg2)
      client sock = sendMany sock [seg1, seg2]

      seg1 = C.pack "This is a "
      seg2 = C.pack "test message."

testRecv :: Test
testRecv = TestCase $ connectedTest client server
    where
      server sock = recv sock 1024 >>= (@=?) testMsg
      client sock = send sock testMsg

testOverFlowRecv :: Test
testOverFlowRecv = TestCase $ connectedTest client server
    where
      server sock = do seg1 <- recv sock (S.length testMsg - 3)
                       seg2 <- recv sock 1024
                       let msg = S.append seg1 seg2
                       testMsg @=? msg

      client sock = send sock testMsg

testRecvFrom :: Test
testRecvFrom = TestCase $ connectedTest client server
    where
      server sock = do (msg, _) <- recvFrom sock 1024
                       testMsg @=? msg

      client sock = send sock testMsg

testOverFlowRecvFrom :: Test
testOverFlowRecvFrom = TestCase $ connectedTest client server
    where
      server sock = do (seg1, _) <- recvFrom sock (S.length testMsg - 3)
                       (seg2, _) <- recvFrom sock 1024
                       let msg = S.append seg1 seg2
                       testMsg @=? msg

      client sock = send sock testMsg

------------------------------------------------------------------------
-- Test helpers

-- | Run a client/server pair and synchronize them so that the server
-- is started before the client and the specified server action is
-- finished before the client closes the connection.
connectedTest :: (Socket -> IO a) -> (Socket -> IO b) -> IO ()
connectedTest clientAct serverAct = do
    barrier <- newEmptyMVar
    forkIO $ server barrier
    client barrier
  where
    server barrier = do
        addr <- inet_addr "127.0.0.1"
        bracket (socket AF_INET Stream defaultProtocol) sClose $ \sock -> do
            setSocketOption sock ReuseAddr 1
            bindSocket sock (SockAddrInet port addr)
            listen sock 1
            serverReady
            (clientSock, _) <- accept sock
            serverAct clientSock
            sClose clientSock
            putMVar barrier ()
      where
        -- | Signal to the client that it can proceed.
        serverReady = putMVar barrier ()

    client barrier = do
        takeMVar barrier
        bracket (socket AF_INET Stream defaultProtocol) sClose $ \sock -> do
            addr <- inet_addr "127.0.0.1"
            connect sock $ SockAddrInet port addr
            clientAct sock
            takeMVar barrier

------------------------------------------------------------------------
-- Test harness

main :: IO ()
main = withSocketsDo $ do
    counts <- runTestTT tests
    when (errors counts + failures counts > 0) exitFailure

tests :: Test
tests = TestList [ TestLabel "testLazySend" testLazySend
                 , TestLabel "testSendAll" testSendMany
                 , TestLabel "testSendMany" testSendMany
                 , TestLabel "testRecv" testRecv
                 , TestLabel "testOverFlowRecv" testOverFlowRecv
                 , TestLabel "testRecvFrom" testRecvFrom
                 , TestLabel "testOverFlowRecvFrom" testOverFlowRecvFrom
                 ]