    {-# OPTIONS_GHC -XDeriveDataTypeable #-}

    module Data.Conduit.Codec.Util
        ( CodecDecodeException(..)
        , encodeI
        , decodeI
        , decodeII
        , encodeII
        ) where

    import Data.Typeable (Typeable)
    import Control.Exception (Exception)
    import Data.ByteString as BS (ByteString, append, null)
    import Data.Conduit (Conduit, MonadThrow, await, monadThrow, yield)
    import Data.Maybe (fromJust)
    import Control.Monad (unless)

    type EncFunc = ByteString -> ByteString
    type EncFuncPart = ByteString -> (ByteString, ByteString)
    type EncFuncFinal = ByteString -> Maybe ByteString
    type DecFunc = ByteString -> Either (ByteString, ByteString) (ByteString, ByteString)
    type DecFuncFinal = ByteString -> Maybe ByteString

    data CodecDecodeException = CodecDecodeException ByteString
        deriving (Typeable, Show)

    instance Exception CodecDecodeException

    encodeI :: (Monad m) => EncFuncPart -> EncFuncFinal -> ByteString -> Conduit ByteString m ByteString
    encodeI enc_part enc_final i = do
        clear <- await
        case clear of
            Nothing -> (yield $ fromJust $ enc_final i) >> return ()
            Just s -> let
                    (a, b) = enc_part (i `append` s)
                in do
                    unless (BS.null a) $ yield a
                    encodeI enc_part enc_final b

    decodeI :: (Monad m, MonadThrow m) => DecFunc -> DecFuncFinal -> ByteString -> Conduit ByteString m ByteString
    decodeI dec_part dec_final i = do
        enc <- await
        case enc of
            Nothing ->
                case dec_final i of
                    Nothing -> monadThrow (CodecDecodeException i)
                    Just s -> yield s >> return ()
            Just s ->
                case dec_part (i `append` s) of
                    Left (a, b) -> do
                        unless (BS.null a) $ yield a
                        monadThrow (CodecDecodeException b)
                    Right (a, b) -> do
                        unless (BS.null a) $ yield a
                        decodeI dec_part dec_final b

    encodeII :: (Monad m) => EncFunc -> Conduit ByteString m ByteString
    encodeII enc = do
        clear <- await
        case clear of
            Nothing -> return ()
            Just s -> do
                yield $ enc s
                encodeII enc

    decodeII :: (Monad m, MonadThrow m) => DecFunc -> ByteString -> Conduit ByteString m ByteString
    decodeII dec i = do
        enc <- await
        case enc of
            Nothing -> if BS.null i
                then return ()
                else monadThrow $ CodecDecodeException i
            Just s -> case (dec $ i `append` s) of
                Left (c, b) -> do
                    unless (BS.null c) $ yield c
                    monadThrow $ CodecDecodeException b
                Right (c, r) -> do
                    unless (BS.null c) $ yield c
                    decodeII dec r