{-# LANGUAGE OverloadedStrings, LambdaCase #-}
module Main where

import Control.Applicative
import Control.Monad (void, forM_, forever, replicateM_)
import Control.Concurrent (forkOS, threadDelay)
import Control.Concurrent.MVar
import Control.Distributed.Process
import Control.Distributed.Process.Node
import Criterion.Types
import Criterion.Measurement as M
import Data.Binary (encode, decode)
import Data.ByteString.Char8 (pack)
import qualified Data.ByteString.Lazy as BSL
import Network.Transport.ZMQ (createTransport, defaultZMQParameters)
import System.Environment
import Text.Printf

pingServer :: Process ()
pingServer = forever $ do
  them <- expect
  send them ()

pingClient :: Int -> ProcessId -> Process ()
pingClient n them = do
  us <- getSelfPid
  replicateM_ n $ send them us >> (expect :: Process ())

initialServer :: Process ()
initialServer = do
  us <- getSelfPid
  liftIO $ BSL.writeFile "pingServer.pid" (encode us)
  pingServer

initialClient :: Int -> Process ()
initialClient n = do
  them <- liftIO $ decode <$> BSL.readFile "pingServer.pid"
  pingClient n them

main :: IO ()
main = getArgs >>= \case
  [] -> defaultBenchmark
  [role, host] -> do
    transport <- createTransport defaultZMQParameters (pack host)
    node <- newLocalNode transport initRemoteTable
    case role of
        "SERVER" -> runProcess node initialServer
        "CLIENT" -> fmap read getLine >>= runProcess node . initialClient
        _        -> error "wrong role"
  _ -> error "either call benchmark with [SERVER|CLIENT] host or without arguments"

defaultBenchmark :: IO ()
defaultBenchmark = do
  -- server
  void . forkOS $ do
    transport <- createTransport defaultZMQParameters "127.0.0.1"
    node <- newLocalNode transport initRemoteTable
    runProcess node $ initialServer
  
  threadDelay 1000000
  e <- newEmptyMVar
  void . forkOS $ do
    putStrLn "pings        time\n---          ---\n"
    forM_ [100,200,600,800,1000,2000,5000,8000,10000] $ \i -> do
        transport <- createTransport defaultZMQParameters "127.0.0.1"
        node <- newLocalNode transport initRemoteTable
        d <- snd <$> M.measure (nfIO $ runProcess node $ initialClient i) 1
        printf "%-8i %10.4f\n" i d
    putMVar e ()
  takeMVar e