{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE FlexibleContexts #-}
module Data.CyclicBuffer where

import Control.Concurrent.STM
import Control.Concurrent.STM.TArray
import Control.Monad
import Control.Applicative
import Data.Word
import Data.Array.MArray
import Data.Maybe

data CyclicBuffer a = CyclicBuffer
    { vwflgs  :: TArray Word32 Bool -- TODO: Use TVar. TArray Word32 (TVar Bool)
                                    --       This would allow updating by external code.
                                    --       The TVar could be returned from dequeue
    , pktq    :: TArray Word32 (Maybe a)
    , seqno   :: TVar Word32
    , qsize   :: Word32
    , buffend :: TVar Word32 -- on incoming, highest packet number handled + 1
    , dropCnt :: TVar Word32
    }

cyclicBufferViewList :: CyclicBuffer a -> STM [(Word32,a)]
cyclicBufferViewList p = do
    let f (n,Nothing) = Nothing
        f (n,Just x)  = Just (n,x)
    catMaybes . map f <$> getAssocs (pktq p)

getCapacity :: Applicative m => CyclicBuffer t -> m Word32
getCapacity (CyclicBuffer { qsize }) = pure qsize

-- | Create a new CyclicBuffer with Overwrite on Wrap.
new :: Word32 -- ^ Capacity of queue.
    -> Word32 -- ^ Initial sequence number.
    -> STM (CyclicBuffer a)
new capacity seqstart = do
    let cap = if capacity `mod` 2 == 0 then capacity else capacity + 1
    q <- newArray (0,cap - 1) Nothing
    flgs <- newArray (0,cap - 1) False
    seqv <- newTVar seqstart
    bufe <- newTVar 0
    dropped <- newTVar 0
    return CyclicBuffer
        { vwflgs  = flgs
        , pktq    = q
        , seqno   = seqv
        , qsize   = cap
        , buffend = bufe
        , dropCnt = dropped
        }

observeOutOfBand :: CyclicBuffer a -> Word32-> STM ()
observeOutOfBand CyclicBuffer { seqno, qsize, buffend } no = do
    low <- readTVar seqno
    let proj = no - low
    -- Ignore packet if out of range.
    when ( proj < qsize) $ do
    modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)


-- | Retry until the next expected packet is enqueued.  Then return it.
dequeue :: CyclicBuffer a -> STM a
dequeue CyclicBuffer { pktq, seqno, qsize } = do
    i0 <- readTVar seqno
    let i = i0 `mod` qsize
    x <- maybe retry return =<< readArray pktq i
    writeArray pktq i Nothing
    modifyTVar' seqno   succ
    return x

-- | Like dequeue, but just marks as viewed rather than removing
markButNotDequeue :: CyclicBuffer a -> STM a
markButNotDequeue CyclicBuffer { vwflgs, pktq, seqno, qsize } = do
    i0 <- readTVar seqno
    let i = i0 `mod` qsize
    x <- maybe retry return =<< readArray pktq i
    writeArray vwflgs i True
    modifyTVar' seqno   succ
    return x

-- | Enqueue a packet.  Packets need not be enqueued in order as long as there
-- is spare capacity in the queue.  If the capacity is exceeded, packets are
-- dropped and the drop count increased accordingly.
enqueue :: CyclicBuffer a -- ^ The cyclic buffer(queue)
        -> Word32        -- ^ Sequence number of the packet.
        -> a             -- ^ The packet.
        -> STM ()
enqueue CyclicBuffer{vwflgs, pktq, seqno, qsize, buffend, dropCnt} no x = do
    low <- readTVar seqno
    let proj = no - low
    let i = no `mod` qsize
    when (proj >= qsize) $ do
        viewed <- readArray vwflgs i
        when (not viewed) $
            modifyTVar' dropCnt (+1)
    writeArray pktq i (Just x)
    writeArray vwflgs i False -- mark as not viewed
    modifyTVar' buffend (\be -> if be - low <= proj then no + 1 else be)
    return ()