{-
Paranoid Pirate Pattern queue in Haskell.
Uses heartbeating to detect crashed or blocked workers.
-}
module Main where

import System.ZMQ4.Monadic
import ZHelpers

import Control.Monad (when, forM_)
import Control.Applicative ((<$>))
import System.IO (hSetEncoding, stdout, utf8)
import Data.ByteString.Char8 (pack, unpack, empty)
import qualified Data.List.NonEmpty as N


type SockID = String

data Worker = Worker {
                       sockID :: SockID
                     , expiry :: Integer
                     } deriving (Show)

heartbeatLiveness = 3
heartbeatInterval_ms = 1000

pppReady = "\001"
pppHeartbeat = "\002"

main :: IO ()
main = 
    runZMQ $ do
        frontend <- socket Router
        bind frontend "tcp://*:5555"
        backend <- socket Router
        bind backend "tcp://*:5556"

        liftIO $ hSetEncoding stdout utf8

        heartbeat_at <- liftIO $ nextHeartbeatTime_ms heartbeatInterval_ms
        pollPeers frontend backend [] heartbeat_at

createWorker :: SockID -> IO Worker
createWorker id = do
    currTime <- currentTime_ms
    let expiry = currTime + heartbeatInterval_ms * heartbeatLiveness
    return (Worker id expiry)

pollPeers :: Socket z Router -> Socket z Router -> [Worker] -> Integer -> ZMQ z ()
pollPeers frontend backend workers heartbeat_at = do
    let toPoll = getPollList workers
    evts <- poll (fromInteger heartbeatInterval_ms) toPoll

    workers' <- getBackend backend frontend evts workers
    workers'' <- getFrontend frontend backend evts workers'

    newHeartbeatAt <- heartbeat backend workers'' heartbeat_at

    workersPurged <- purge workers''

    pollPeers frontend backend workersPurged newHeartbeatAt
    
  where getPollList [] = [Sock backend [In] Nothing]
        getPollList _  = [Sock backend [In] Nothing, Sock frontend [In] Nothing]

        getBackend :: Socket z Router -> Socket z Router ->
                      [[Event]] -> [Worker] -> ZMQ z ([Worker])
        getBackend backend frontend evts workers =
            if (In `elem` (evts !! 0))
            then do
                frame <- receiveMulti backend
                let wkrID = frame !! 0
                    msg = frame !! 1

                if ((length frame) == 2) -- PPP message
                then when (unpack msg `notElem` [pppReady, pppHeartbeat]) $ do
                    liftIO $ putStrLn $ "E: Invalid message from worker " ++ (unpack msg)
                else do -- Route the message to the client
                    liftIO $ putStrLn "I: Sending normal message to client"
                    let id = frame !! 1
                        msg = frame !! 3
                    send frontend [SendMore] id
                    send frontend [SendMore] empty
                    send frontend [] msg

                newWorker <- liftIO $ createWorker $ unpack wkrID
                return $ workers ++ [newWorker]
            else return workers

        getFrontend :: Socket z Router -> Socket z Router ->
                       [[Event]] -> [Worker] -> ZMQ z ([Worker])
        getFrontend frontend backend evts workers =
            if (length evts > 1 && In `elem` (evts !! 1))
            then do -- Route message to workers
                frame <- receiveMulti frontend
                
                let wkrID = sockID . head $ workers
                send backend [SendMore] (pack wkrID)
                send backend [SendMore] empty
                sendMulti backend (N.fromList frame)
                return $ tail workers
            else return workers
        
        heartbeat :: Socket z Router -> [Worker] -> Integer -> ZMQ z Integer
        heartbeat backend workers heartbeat_at = do
            currTime <- liftIO currentTime_ms
            if (currTime >= heartbeat_at) 
            then do
                forM_ workers (\worker -> do
                    send backend [SendMore] (pack $ sockID worker)
                    send backend [SendMore] empty
                    send backend [] (pack pppHeartbeat)
                    liftIO $ putStrLn $ "I: sending heartbeat to '" ++ (sockID worker) ++ "'")
                liftIO $ nextHeartbeatTime_ms heartbeatInterval_ms
            else return heartbeat_at

        purge :: [Worker] -> ZMQ z ([Worker])
        purge workers = do
            currTime <- liftIO currentTime_ms
            return $ filter (\wkr -> expiry wkr > currTime) workers
