module StreamPatch.Stream where

import           GHC.Natural
import           Data.Kind
import qualified Data.ByteString            as BS
import qualified Data.ByteString.Builder    as BB
import           Control.Monad.State
import           Control.Monad.Reader
import           System.IO                  ( Handle, SeekMode(..), hSeek )

import qualified Data.List                  as List

-- TODO also require reporting cursor position (for error reporting)
class Monad m => MonadFwdInplaceStream m where
    type Chunk m :: Type

    -- | Read a number of bytes without advancing the cursor.
    readahead :: Natural -> m (Chunk m)

    -- | Advance cursor without reading.
    advance :: Natural -> m ()

    -- | Insert bytes into the stream at the cursor position, overwriting
    --   existing bytes.
    overwrite :: Chunk m -> m ()

-- TODO Need MonoTraversable to define for Text, ByteString etc easily. Bleh. I
-- think Snoyman's advice is to reimplement. Also bleh.
instance Monad m => MonadFwdInplaceStream (StateT ([a], [a]) m) where
    type Chunk (StateT ([a], [a]) m) = [a]
    readahead :: Natural -> StateT ([a], [a]) m (Chunk (StateT ([a], [a]) m))
readahead Natural
n = Int -> [a] -> [a]
forall a. Int -> [a] -> [a]
List.take (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n) ([a] -> [a]) -> StateT ([a], [a]) m [a] -> StateT ([a], [a]) m [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (([a], [a]) -> [a]) -> StateT ([a], [a]) m [a]
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets ([a], [a]) -> [a]
forall a b. (a, b) -> a
fst
    advance :: Natural -> StateT ([a], [a]) m ()
advance Natural
n = do
        ([a]
src, [a]
out) <- StateT ([a], [a]) m ([a], [a])
forall s (m :: * -> *). MonadState s m => m s
get
        let ([a]
bs, [a]
src') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
List.splitAt (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n) [a]
src
        ([a], [a]) -> StateT ([a], [a]) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([a]
src', [a]
out [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
bs)
    overwrite :: Chunk (StateT ([a], [a]) m) -> StateT ([a], [a]) m ()
overwrite Chunk (StateT ([a], [a]) m)
bs = do
        ([a]
src, [a]
out) <- StateT ([a], [a]) m ([a], [a])
forall s (m :: * -> *). MonadState s m => m s
get
        let ([a]
_, [a]
src') = Int -> [a] -> ([a], [a])
forall a. Int -> [a] -> ([a], [a])
List.splitAt ([a] -> Int
forall (t :: * -> *) a. Foldable t => t a -> Int
List.length [a]
Chunk (StateT ([a], [a]) m)
bs) [a]
src
        ([a], [a]) -> StateT ([a], [a]) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put ([a]
src', [a]
out [a] -> [a] -> [a]
forall a. Semigroup a => a -> a -> a
<> [a]
Chunk (StateT ([a], [a]) m)
bs)

instance MonadIO m => MonadFwdInplaceStream (ReaderT Handle m) where
    type Chunk (ReaderT Handle m) = BS.ByteString
    readahead :: Natural -> ReaderT Handle m (Chunk (ReaderT Handle m))
readahead Natural
n = do
        Handle
hdl <- ReaderT Handle m Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
        ByteString
bs <- IO ByteString -> ReaderT Handle m ByteString
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO ByteString -> ReaderT Handle m ByteString)
-> IO ByteString -> ReaderT Handle m ByteString
forall a b. (a -> b) -> a -> b
$ Handle -> Int -> IO ByteString
BS.hGet Handle
hdl (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n)
        IO () -> ReaderT Handle m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT Handle m ()) -> IO () -> ReaderT Handle m ()
forall a b. (a -> b) -> a -> b
$ Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
hdl SeekMode
RelativeSeek (- Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n)
        ByteString -> ReaderT Handle m ByteString
forall (m :: * -> *) a. Monad m => a -> m a
return ByteString
bs
    advance :: Natural -> ReaderT Handle m ()
advance Natural
n = do
        Handle
hdl <- ReaderT Handle m Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
        IO () -> ReaderT Handle m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT Handle m ()) -> IO () -> ReaderT Handle m ()
forall a b. (a -> b) -> a -> b
$ Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
hdl SeekMode
RelativeSeek (Natural -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n)
    overwrite :: Chunk (ReaderT Handle m) -> ReaderT Handle m ()
overwrite Chunk (ReaderT Handle m)
bs = do
        Handle
hdl <- ReaderT Handle m Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
        IO () -> ReaderT Handle m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT Handle m ()) -> IO () -> ReaderT Handle m ()
forall a b. (a -> b) -> a -> b
$ Handle -> ByteString -> IO ()
BS.hPut Handle
hdl ByteString
Chunk (ReaderT Handle m)
bs

instance Monad m => MonadFwdInplaceStream (StateT (BS.ByteString, BB.Builder) m) where
    type Chunk (StateT (BS.ByteString, BB.Builder) m) = BS.ByteString
    readahead :: Natural
-> StateT
     (ByteString, Builder) m (Chunk (StateT (ByteString, Builder) m))
readahead Natural
n = Int -> ByteString -> ByteString
BS.take (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n) (ByteString -> ByteString)
-> StateT (ByteString, Builder) m ByteString
-> StateT (ByteString, Builder) m ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ((ByteString, Builder) -> ByteString)
-> StateT (ByteString, Builder) m ByteString
forall s (m :: * -> *) a. MonadState s m => (s -> a) -> m a
gets (ByteString, Builder) -> ByteString
forall a b. (a, b) -> a
fst
    advance :: Natural -> StateT (ByteString, Builder) m ()
advance Natural
n = do
        (ByteString
src, Builder
out) <- StateT (ByteString, Builder) m (ByteString, Builder)
forall s (m :: * -> *). MonadState s m => m s
get
        let (ByteString
bs, ByteString
src') = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt (Natural -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Natural
n) ByteString
src
        (ByteString, Builder) -> StateT (ByteString, Builder) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (ByteString
src', Builder
out Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
BB.byteString ByteString
bs)
    overwrite :: Chunk (StateT (ByteString, Builder) m)
-> StateT (ByteString, Builder) m ()
overwrite Chunk (StateT (ByteString, Builder) m)
bs = do
        (ByteString
src, Builder
out) <- StateT (ByteString, Builder) m (ByteString, Builder)
forall s (m :: * -> *). MonadState s m => m s
get
        let (ByteString
_, ByteString
src') = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt (ByteString -> Int
BS.length ByteString
Chunk (StateT (ByteString, Builder) m)
bs) ByteString
src
        (ByteString, Builder) -> StateT (ByteString, Builder) m ()
forall s (m :: * -> *). MonadState s m => s -> m ()
put (ByteString
src', Builder
out Builder -> Builder -> Builder
forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
BB.byteString ByteString
Chunk (StateT (ByteString, Builder) m)
bs)

-- | A forward stream, but backward too.
class MonadFwdInplaceStream m => MonadCursorInplaceStream m where
    -- | Move cursor.
    move :: Integer -> m ()

instance MonadIO m => MonadCursorInplaceStream (ReaderT Handle m) where
    move :: Integer -> ReaderT Handle m ()
move Integer
n = do
        Handle
hdl <- ReaderT Handle m Handle
forall r (m :: * -> *). MonadReader r m => m r
ask
        IO () -> ReaderT Handle m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> ReaderT Handle m ()) -> IO () -> ReaderT Handle m ()
forall a b. (a -> b) -> a -> b
$ Handle -> SeekMode -> Integer -> IO ()
hSeek Handle
hdl SeekMode
RelativeSeek (Integer -> Integer
forall a b. (Integral a, Num b) => a -> b
fromIntegral Integer
n)

class MonadFwdInplaceStream m => MonadFwdStream m where
    insert :: Chunk m -> m ()