{-# LANGUAGE TupleSections #-}
module Ros.Topic.Stats (sendMessageStat, recvMessageStat, statSnapshot,
                        StatMap, SubStats(..), PubStats(..)) where
import Control.Applicative ((<$>))
import Control.Concurrent.STM.TVar
import Control.Concurrent.STM
import Data.Map (Map)
import qualified Data.Map as M
import Ros.Internal.RosTypes

data SubStats = SubStats { bytesReceived  :: !Int
                         , subConnected   :: !Bool }

data PubStats = PubStats { bytesSent      :: !Int
                         , numSent        :: !Int
                         , pubConnected   :: !Bool }


-- |A transactional data store for tracking all the connections for a
-- particular topic.
type StatMap a = TVar (Map URI (TVar a))

statSnapshot :: (URI, TVar a) -> STM (URI, a)
statSnapshot (uri, stat) = (uri,) <$> readTVar stat

-- |Record the fact that we've sent a message of the given number of
-- bytes to the given URI. If the number of bytes is negative, the
-- connection is marked is disconnected.
sendMessageStat :: StatMap PubStats -> URI -> Int -> IO ()
sendMessageStat tm uri numBytes = 
    atomically $ do m <- readTVar tm
                    let conn = numBytes >= 0
                        nb = max 0 numBytes
                        nm = if conn then 1 else 0
                    case M.lookup uri m of
                      Nothing -> do stats <- newTVar $ PubStats nb nm conn
                                    writeTVar tm (M.insert uri stats m)
                      Just ts -> do PubStats nb' nm' _ <- readTVar ts
                                    let stats = PubStats (nb' + nb)
                                                         (nm' + nm)
                                                         conn
                                    writeTVar ts stats

-- |Record the fact that we've received a message of the given number
-- of bytes from the given URI. If the number of bytes is negative,
-- the connection is marked as disconnected.
recvMessageStat :: StatMap SubStats -> URI -> Int -> IO ()
recvMessageStat tm uri numBytes = 
    atomically $ do m <- readTVar tm
                    let conn = numBytes >= 0
                        nb = max 0 numBytes
                    case M.lookup uri m of
                      Nothing -> do stats <- newTVar $ SubStats nb conn
                                    writeTVar tm (M.insert uri stats m)
                      Just ts -> do SubStats nb' _ <- readTVar ts
                                    let stats = SubStats (nb' + nb) conn
                                    writeTVar ts stats