{-# LANGUAGE BangPatterns, ScopedTypeVariables, FlexibleContexts #-}

module Data.MutableIter.IOBuffer (
  IOBuffer
  ,createIOBuffer
  ,null
  ,empty
  ,copyBuffer
  ,append
  ,length
  ,pop
  ,lookAtHead
  ,drop
  ,dropWhile
  ,take
  ,splitAt
  ,mapBuffer
  ,mapAccumBuffer
  ,foldl'
  ,hopfoldl'
  ,hopfoldM
  ,decimate
  ,castBuffer
  ,freeze
  ,thaw
  ,hPut
  ,unsafeToForeignPtr
)

where

import Prelude hiding (length, take, drop, null, splitAt, dropWhile)

import qualified Data.Vector.Generic as G
import qualified Data.Iteratee as I
import Control.Monad
import Control.Monad.CatchIO

import Foreign
import System.IO

import System.IO.Unsafe (unsafePerformIO)


-- |A mutable buffer to hold storable elements.  This data type supports
-- memory recycling.
data IOBuffer r el = IOBuffer {-# UNPACK #-} !Int
                              {-# UNPACK #-} !(ForeignPtr Int)
                              {-# UNPACK #-} !(ForeignPtr el)

instance Storable el => I.NullPoint (IOBuffer r el) where
  empty = empty

fpPeek :: Storable el => ForeignPtr el -> IO el
fpPeek = flip withForeignPtr peek

fpPoke :: Storable el => ForeignPtr el -> el -> IO ()
fpPoke fp el = withForeignPtr fp (flip poke el)

withBuf :: IOBuffer r el -> (Int -> Ptr Int -> Ptr el -> IO a) -> IO a
withBuf (IOBuffer l ofp ofb) f = withForeignPtr ofp (\op ->
  withForeignPtr ofb (\ob -> f l op ob))
{-# INLINE withBuf #-}

newFp :: Storable a => a -> IO (ForeignPtr a)
newFp a = mallocForeignPtr >>= \fp -> fpPoke fp a >> return fp

-- |Create a buffer from a length and data array.
createIOBuffer :: (Storable el) =>
  Int
  -> ForeignPtr Int
  -> ForeignPtr el
  -> IOBuffer r el
createIOBuffer len op buf = IOBuffer len op buf

-- |Empty buffer.
empty :: Storable el => IOBuffer r el
empty = IOBuffer 0 nullForeignPtr nullForeignPtr

nullForeignPtr = unsafePerformIO (newForeignPtr_ nullPtr)

-- |Check if the buffer is empty.
null :: IOBuffer r el -> IO Bool
null (IOBuffer 0 _  _) = return True
null buf = withBuf buf $ \l po _ -> liftM (>= l) $ peek po
{-# INLINE null #-}

-- |IOBuffer length.
length :: IOBuffer r el -> IO Int
length buf = withBuf buf $ \l po _ -> liftM (l -) $ peek po
{-# INLINE length #-}

-- |Retrieve the front element from the buffer and advance the internal pointer.
-- It is an error to call this on an empty buffer.
pop :: (Storable el) => IOBuffer r el -> IO el
pop (IOBuffer 0 _  _ ) = error "Can't pop head off of empty buffer"
pop buf = withBuf buf $ \l po pb -> do
  off <- peek po
  if off >= l then error "Can't pop head from empty buffer"
              else poke po (off+1) >> peek (pb `advancePtr` off)
{-# INLINE pop #-}

-- |Retrieve the first element, if it exists.
-- This function does not advance the buffer pointer.
lookAtHead :: (Storable el)  =>
  IOBuffer r el -> IO (Maybe el)
lookAtHead (IOBuffer 0 _  _ ) = return Nothing
lookAtHead buf = withBuf buf $ \l po pb -> do
  off <- peek po
  if off >= l then return Nothing
              else liftM Just $ peek (pb `advancePtr` off)
{-# INLINE lookAtHead #-}

-- |Drop n elements from the front of the buffer.
-- if the buffer has fewer elements, all are dropped.
drop :: Int -> IOBuffer r el -> IO ()
drop n_drop buf = withBuf buf $ \l po pb -> do
  off <- peek po
  poke po (min l (off+n_drop))
{-# INLINE drop #-}

dropWhile :: (Storable el) =>
  (el -> Bool)
  -> IOBuffer r el
  -> IO ()
dropWhile pred buf = withBuf buf $ \l po pb -> do
  off <- peek po
  let len = l-off
  let go cnt p | cnt < len = do
          this <- peek p
          if pred this then go (succ cnt) (p `advancePtr` 1) else return cnt
  let go cnt _ = return cnt --off the end of the buffer
  n <- go 0 (pb `advancePtr` off)
  poke po n
{-# INLINE dropWhile #-}

-- |Create a new buffer from the first n elements, sharing data.
-- This function advances the pointer of the original buffer.
take :: (Storable el) =>
  IOBuffer r el
  -> Int
  -> IO (IOBuffer r el)
take (IOBuffer 0 _  _ ) _      = return empty
take buf@(IOBuffer _ fpo fpb) n_take = withBuf buf $ \l po pb -> do
  off <- peek po
  po' <- newFp off
  poke po $ min l (off+n_take)
  return $ IOBuffer l po' fpb

-- |Split one buffer to two, sharing storage.
splitAt :: (Storable el) =>
  IOBuffer r el
  -> Int
  -> IO (IOBuffer r el, IOBuffer r el)
splitAt (IOBuffer 0 _ _) _ = return (empty, empty)
splitAt buf              0 = return (empty, buf)
splitAt buf@(IOBuffer l fpo fpb) n
  | n>0 && n <= l = withForeignPtr fpo $ \po -> withForeignPtr fpb $ \pb -> do
    let ib1 = IOBuffer n fpo fpb
    off <- peek po
    po2 <- newFp (off+n)
    return (ib1, IOBuffer l po2 fpb)
  | True = return (buf, empty)

-- |Copy data from one buffer to another.
copyBuffer :: (Storable el) =>
  IOBuffer r el -> IO (IOBuffer r el)
copyBuffer (IOBuffer 0 _ _) = return empty
copyBuffer buf = withBuf buf $ \l po pb -> do
  off <- peek po
  let len' = l-off
  if len' > 0 then do
      po' <- newFp 0
      pb' <- mallocForeignPtrArray len'
      withForeignPtr pb' $ \p -> copyArray p (pb `advancePtr` off) len'
      return $ IOBuffer len' po' pb'
    else return empty

-- |Append two buffers.  Copies data from both into a new buffer.
append :: (Storable el) =>
  IOBuffer r el
  -> IOBuffer r el
  -> IO (IOBuffer r el)
append ib1 (IOBuffer 0 _ _) = copyBuffer ib1
append (IOBuffer 0 _ _) ib2 = copyBuffer ib2
append ib1 ib2 = withBuf ib1 (\l1 po1 pb1 -> withBuf ib2 (\l2 po2 pb2 -> do
  len1 <- length ib1
  len2 <- length ib2
  let len = len1 + len2
  off1 <- peek po1
  off2 <- peek po2
  po <- newFp 0
  pb <- mallocForeignPtrArray len
  withForeignPtr pb $ \p -> copyArray p (pb1 `advancePtr` off1) len1
  withForeignPtr pb $ \p -> copyArray (p `advancePtr` len1)
                            (pb2 `advancePtr` off2) len2
  return $ IOBuffer len po pb))

-- |Safely convert an IOBuffer to a Vector.
freeze :: (Storable el, G.Vector v el, G.Vector v Int) =>
  IOBuffer r el
  -> IO (v el)
freeze (IOBuffer 0 _ _) = return G.empty
freeze buf = withBuf buf $ \l po pb -> do
  len <- length buf
  off <- peek po
  let vIx = G.enumFromN 0 len
  vRes <- G.forM vIx (\ix -> peekElemOff (pb `advancePtr` off) ix)
  return vRes
{-# INLINE freeze #-}

-- |Safely convert a Vector to an IOBuffer
thaw :: (Storable el, G.Vector v el) => v el -> IO (IOBuffer r el)
thaw vec = do
  offp <- newFp 0
  bufp <- mallocForeignPtrArray (G.length vec)
  withForeignPtr bufp $ \p ->
    G.foldM' (\ix el -> pokeElemOff p ix el >> return (ix+1)) 0 vec
  return $ createIOBuffer (G.length vec) offp bufp
{-# INLINE thaw #-}

-- |Write out the contents of the IOBuffer to a handle.  This operation
-- drains the buffer.
hPut :: forall r el. (Storable el) => Handle -> IOBuffer r el -> IO ()
hPut h (IOBuffer 0 _ _)   = return ()
hPut h buf = withBuf buf $ \l po pb -> do
  off <- peek po
  (hPutBuf h) (pb `advancePtr` off) (bytemult * (l-off))
  poke po l
  where
    bytemult = sizeOf (undefined :: el)

-- |copy data from one buffer to another with the specified map function.
-- this operation drains the original buffer.
mapBuffer :: (Storable el, Storable el') =>
  (el -> el')
  -> ForeignPtr Int
  -> ForeignPtr el'
  -> IOBuffer r el
  -> IO (IOBuffer r el')
mapBuffer f dp datap (IOBuffer 0 _ _) = return empty
mapBuffer f fdp fdatap buf = withForeignPtr fdp $ \dp ->
  withForeignPtr fdatap $ \datap -> withBuf buf $ \l po pb ->
    let go !cnt ip' op' = when (cnt>0) $ do
          poke op' . f =<< peek ip'
          go (cnt-1) (ip' `advancePtr` 1) (op' `advancePtr` 1)
    in do
         off <- peek po
         go (l-off) (pb `advancePtr` off) datap
         poke dp 0
         poke po l
         return $ createIOBuffer (l-off) fdp fdatap
{-# INLINE mapBuffer #-}

mapAccumBuffer :: (Storable el, Storable el') =>
  (acc -> el -> (acc,el'))
  -> ForeignPtr Int
  -> ForeignPtr el'
  -> acc
  -> IOBuffer r el
  -> IO (acc, IOBuffer r el')
mapAccumBuffer f fdp fdatap acc (IOBuffer 0 _ _) = return (acc, empty)
mapAccumBuffer f fdp fdatap acc buf = withForeignPtr fdp $ \dp ->
  withForeignPtr fdatap $ \datap -> withBuf buf $ \l po pb ->
    let go !acc' cnt !ip' !op'
          | cnt == 0 = return acc'
          | True = do
            (acc2, el) <- liftM (f acc') $ peek ip'
            poke op' el
            go acc2 (cnt-1) (ip' `advancePtr` 1) (op' `advancePtr` 1)
    in do
         off  <- peek po
         acc' <- go acc (l-off) (pb `advancePtr` off) datap
         poke dp 0
         poke po l
         return $ (acc', createIOBuffer (l-off) fdp fdatap)
{-# INLINE mapAccumBuffer #-}


-- |Cast a buffer to a different type.  Any extra data is truncated.
-- This is not safe unless the buffer offset is 0.
castBuffer :: forall r m el el'. (Storable el, Storable el') =>
  IOBuffer r el -> IO (IOBuffer r el')
castBuffer (IOBuffer 0 _ _) = return empty
castBuffer (IOBuffer l po pb) = do
  off <- fpPeek po
  when (off /= 0) (error "castBuffer called with non-zero offset")
  return $ IOBuffer l' po (castForeignPtr pb)
  where
    l' = (l * sizeOf (undefined :: el)) `div` sizeOf (undefined :: el')

foldl' :: (Storable b) => (a -> b -> a) -> a -> IOBuffer r b -> IO a
foldl' f acc (IOBuffer 0 _ _) = return acc
foldl' f i0 buf = withBuf buf $ \l po pb ->
  let go !n !acc p | n>0 = do
        el <- peek p
        go (n-1) (f acc el) (p `advancePtr` 1)
      go _ acc _ = return acc
  in do
    off <- peek po
    go (l-off) i0 (pb `advancePtr` off)
{-# INLINE foldl' #-}

hopfoldl' :: (Storable b) => Int -> (a -> b -> a) -> a -> IOBuffer r b -> IO a
hopfoldl' hop f acc (IOBuffer 0 _ _) = return acc
hopfoldl' hop f i0 buf = withBuf buf $ \l po pb ->
  let go !n !acc p | n>0 = do
        el <- peek p
        go (n-1) (f acc el) (p `advancePtr` hop)
      go _ acc _ = return acc
  in do
    off <- peek po
    go (l-off `div` hop) i0 (pb `advancePtr` off)
{-# INLINE hopfoldl' #-}

hopfoldM :: (Storable b) => Int -> (a -> b -> IO a) -> a -> IOBuffer r b -> IO a
hopfoldM hop f acc (IOBuffer 0 _ _) = return acc
hopfoldM hop f i0 buf = withBuf buf $ \l po pb ->
  let go !n !acc p | n>0 = do
        el <- peek p
        acc' <- f acc el
        go (n-1) acc' (p `advancePtr` hop)
      go _ acc _ = return acc
  in do
    off <- peek po
    go ((l-off) `div` hop) i0 (pb `advancePtr` off)
{-# INLINE hopfoldM #-}

-- | Create a new buffer of every 'nth' element.  The original buffer is
--   not altered.
decimate :: Storable b => Int -> IOBuffer r b -> IO (IOBuffer r b)
decimate _skp (IOBuffer 0 _ _) = return empty
decimate skip buf = do
  ln <- length buf
  let newlen = ln `div` skip
  bufp <- mallocForeignPtrArray newlen
  withForeignPtr bufp $ \p' ->
    hopfoldM skip (\p el -> poke p el >> return (p `advancePtr` 1)) p' buf
  op <- newFp 0
  return $ createIOBuffer newlen op bufp

unsafeToForeignPtr :: IOBuffer r el -> (Int, ForeignPtr Int, ForeignPtr el)
unsafeToForeignPtr (IOBuffer l po pb) = (l,po,pb)