module Network.Control.Rate (
    -- * Rate control
    Rate,
    newRate,
    getRate,
    addRate,
) where

import Data.IORef
import Data.UnixTime

-- | Type for rating.
newtype Rate = Rate (IORef Counter)

data Counter = Counter Int UnixTime

-- | Creating a new 'Rate'.
newRate :: IO Rate
newRate :: IO Rate
newRate = do
    Counter
cntr <- Int -> UnixTime -> Counter
Counter Int
0 (UnixTime -> Counter) -> IO UnixTime -> IO Counter
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO UnixTime
getUnixTime
    IORef Counter -> Rate
Rate (IORef Counter -> Rate) -> IO (IORef Counter) -> IO Rate
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Counter -> IO (IORef Counter)
forall a. a -> IO (IORef a)
newIORef Counter
cntr

-- | Getting the current rate.
-- If one or more seconds have passed since the previous call, the
-- counter is re-initialized with 1 and it is returned.  Otherwise,
-- incremented counter number is returned.
getRate :: Rate -> IO Int
getRate :: Rate -> IO Int
getRate Rate
r = Rate -> Int -> IO Int
addRate Rate
r Int
1

-- | Getting the current rate.
-- If one or more seconds have passed since the previous call, the
-- counter is re-initialized with the second argument and it is
-- returned.  Otherwise, increased counter number is returned.
addRate :: Rate -> Int -> IO Int
addRate :: Rate -> Int -> IO Int
addRate (Rate IORef Counter
ref) Int
x = do
    Counter Int
n UnixTime
beg <- IORef Counter -> IO Counter
forall a. IORef a -> IO a
readIORef IORef Counter
ref
    UnixTime
cur <- IO UnixTime
getUnixTime
    if (UnixTime
cur UnixTime -> UnixTime -> UnixDiffTime
`diffUnixTime` UnixTime
beg) UnixDiffTime -> UnixDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
> UnixDiffTime
1
        then do
            let n' :: Int
n' = Int
x
            IORef Counter -> Counter -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Counter
ref (Counter -> IO ()) -> Counter -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> UnixTime -> Counter
Counter Int
n' UnixTime
cur
            Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n'
        else do
            let n' :: Int
n' = Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
x
            IORef Counter -> Counter -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef IORef Counter
ref (Counter -> IO ()) -> Counter -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> UnixTime -> Counter
Counter Int
n' UnixTime
beg
            Int -> IO Int
forall a. a -> IO a
forall (m :: * -> *) a. Monad m => a -> m a
return Int
n'