module ConcurrentBuffer
(
  Buffer,
  new,
  push,
  pushBytes,
  pushStorable,
  pull,
  pullBytes,
  pullStorable,
  getSpace,
  getBytes,
)
where

import ConcurrentBuffer.Prelude hiding (State, Buffer)
import qualified ConcurrentBuffer.PtrIO as A
import qualified Data.ByteString.Internal as C


data Buffer =
  {-|
  * Buffer pointer
  * Start offset
  * End offset
  * Capacity
  -}
  Buffer
    {-# UNPACK #-} !(TVar (ForeignPtr Word8))
    {-# UNPACK #-} !(TVar Int)
    {-# UNPACK #-} !(TVar Int)
    {-# UNPACK #-} !(TVar Int)
    {-# UNPACK #-} !(TVar Bool)
    {-# UNPACK #-} !(TVar Bool)
    {-# UNPACK #-} !(TVar Bool)


{-|
Create a new buffer of the specified initial capacity.
-}
new :: Int -> IO Buffer
new capacity =
  do
    fptr <- mallocForeignPtrBytes capacity
    atomically $ do
      fptrVar <- newTVar fptr
      startVar <- newTVar 0
      endVar <- newTVar 0
      capVar <- newTVar capacity
      notPullingVar <- newTVar True
      notPushingVar <- newTVar True
      notAligningVar <- newTVar True
      return (Buffer fptrVar startVar endVar capVar notPullingVar notPushingVar notAligningVar)

{-|
Prepares the buffer to be filled with at maximum the specified amount of bytes,
then uses the pointer-action to populate it.
It is your responsibility to ensure that the action does not exceed the space limit.

The pointer-action returns the amount of bytes it actually writes to the buffer.
That amount then is used to move the buffer's cursor accordingly.
It can also produce some @result@, which will then be emitted by @push@.

It also aligns or grows the buffer if required.
-}
push :: Buffer -> Int -> (Ptr Word8 -> IO (Int, result)) -> IO result
push (Buffer fptrVar startVar endVar capVar notPullingVar notPushingVar notAligningVar) space ptrIO =
  join $ atomically $ do
    notPushing <- readTVar notPushingVar
    guard notPushing
    writeTVar notPushingVar False
    fptr <- readTVar fptrVar
    start <- readTVar startVar
    end <- readTVar endVar
    capacity <- readTVar capVar
    let
      !remainingSpace = capacity - end
      !capacityDelta = space - remainingSpace
      !occupiedSpace = end - start
    if capacityDelta <= 0 -- Doesn't need more space?
      then
        return $ do
          (!pushedSpace, output) <- withForeignPtr fptr $ \ptr -> ptrIO (plusPtr ptr end)
          atomically $ do
            writeTVar endVar $! end + pushedSpace
            writeTVar notPushingVar True
          return output
      else if capacityDelta > start -- Needs growing?
        then
          -- Grow
          return $ do
            let !newCapacity = occupiedSpace + space
            newFPtr <- mallocForeignPtrBytes newCapacity
            (!pushedSpace, output) <- withForeignPtr newFPtr $ \newPtr -> do
              withForeignPtr fptr $ \ptr -> A.memcpy newPtr (plusPtr ptr start) (fromIntegral occupiedSpace)
              ptrIO (plusPtr newPtr occupiedSpace)
            let !newEnd = occupiedSpace + pushedSpace
            atomically $ do
              divergedStart <- readTVar startVar
              let !newStart = divergedStart - start
              writeTVar fptrVar newFPtr
              writeTVar startVar newStart
              writeTVar endVar newEnd
              writeTVar capVar newCapacity
              writeTVar notPushingVar True
            return output
        else
          -- Align
          do
            notPulling <- readTVar notPullingVar
            guard notPulling
            writeTVar notAligningVar False
            return $ do
              (!pushedSpace, output) <- withForeignPtr fptr $ \ptr -> do
                A.memmove ptr (plusPtr ptr start) (fromIntegral occupiedSpace)
                atomically $ do
                  writeTVar startVar 0
                  writeTVar endVar occupiedSpace
                  writeTVar notAligningVar True
                ptrIO (plusPtr ptr occupiedSpace)
              atomically $ do
                writeTVar endVar $! occupiedSpace + pushedSpace
                writeTVar notPushingVar True
              return output

{-|
Pulls the specified amount of bytes from the buffer using the provided pointer-action,
freeing the buffer from the pulled bytes afterwards.

In case the buffer does not contain enough bytes yet, it will block waiting.
-}
pull :: Buffer -> Int -> (Ptr Word8 -> IO result) -> IO result
pull (Buffer fptrVar startVar endVar capVar notPullingVar notPushingVar notAligningVar) amount ptrIO =
  join $ atomically $ do
    notPulling <- readTVar notPullingVar
    guard notPulling
    fptr <- readTVar fptrVar
    start <- readTVar startVar
    end <- readTVar endVar
    guard (amount <= end - start)
    notAligning <- readTVar notAligningVar
    guard notAligning
    writeTVar notPullingVar False
    return $ do
      pulled <- withForeignPtr fptr $ \ptr -> ptrIO (plusPtr ptr start)
      atomically $ do
        start <- readTVar startVar
        writeTVar startVar $! start + amount
        writeTVar notPullingVar True
      return pulled

{-|
Push a byte array into the buffer.
-}
{-# INLINE pushBytes #-}
pushBytes :: Buffer -> ByteString -> IO ()
pushBytes buffer (C.PS bytesFPtr offset length) =
  push buffer length $ \ptr ->
  withForeignPtr bytesFPtr $ \bytesPtr ->
  C.memcpy ptr (plusPtr bytesPtr offset) length $> (length, ())

{-|
Pulls the specified amount of bytes.
-}
{-# INLINE pullBytes #-}
pullBytes :: Buffer -> Int -> IO ByteString
pullBytes buffer amount =
  pull buffer amount (\ptr -> C.create amount (\destPtr -> C.memcpy destPtr ptr amount))

{-|
Push a storable value into the buffer.
-}
{-# INLINE pushStorable #-}
pushStorable :: (Storable storable) => Buffer -> storable -> IO ()
pushStorable buffer storable =
  push buffer amount (\ptr -> poke (castPtr ptr) storable $> (amount, ()))
  where
    amount = sizeOf storable

{-|
Pulls a storable value.
-}
{-# INLINE pullStorable #-}
pullStorable :: (Storable storable) => Buffer -> IO storable
pullStorable buffer =
  result
  where
    result =
      pull buffer amount (\ptr -> peek (castPtr ptr))
      where
        amount =
          sizeOf ((undefined :: IO a -> a) result)

{-|
Get how much space is occupied by the buffer's data.
-}
{-# INLINE getSpace #-}
getSpace :: Buffer -> IO Int
getSpace (Buffer fptrVar startVar endVar capVar notPullingVar notPushingVar notAligningVar) =
  atomically $ do
    end <- readTVar endVar
    start <- readTVar startVar
    return $! end - start

{-|
Create a bytestring representation without modifying the buffer.
-}
{-# INLINE getBytes #-}
getBytes :: Buffer -> IO ByteString
getBytes (Buffer fptrVar startVar endVar capVar notPullingVar notPushingVar notAligningVar) =
  join $ atomically $ do
    fptr <- readTVar fptrVar
    end <- readTVar endVar
    start <- readTVar startVar
    let size = end - start
    return $ withForeignPtr fptr $ \ptr -> C.create size $ \destPtr -> C.memcpy destPtr ptr size