{-# LANGUAGE DeriveDataTypeable, RankNTypes #-}

-- |

-- Copyright: 2011 Michael Snoyman, 2010-2011 John Millikin

-- License: MIT

--

-- Handle streams of text.

--

-- Parts of this code were taken from enumerator and adapted for conduits.

--

-- For many purposes, it's recommended to use the conduit-combinators library,

-- which provides a more complete set of functions.

module Data.Conduit.Text

    (

    -- * Text codecs

      Codec

    , encode

    , decode

    , utf8

    , utf16_le

    , utf16_be

    , utf32_le

    , utf32_be

    , ascii

    , iso8859_1

    , lines

    , linesBounded

    , TextException (..)

    , takeWhile

    , dropWhile

    , take

    , drop

    , foldLines

    , withLine

    , Data.Conduit.Text.decodeUtf8

    , decodeUtf8Lenient

    , encodeUtf8

    , detectUtf

    ) where

import           Prelude hiding (head, drop, takeWhile, lines, zip, zip3, zipWith, zipWith3, take, dropWhile)

import qualified Control.Exception as Exc

import qualified Data.ByteString as B

import qualified Data.ByteString.Char8 as B8

import           Data.Char (ord)

import qualified Data.Text as T

import qualified Data.Text.Encoding as TE

import           Data.Word (Word8)

import           Data.Typeable (Typeable)

import Data.Conduit

import qualified Data.Conduit.List as CL

import Control.Monad.Trans.Class (lift)

import Control.Monad.Trans.Resource (MonadThrow, monadThrow)

import Control.Monad (unless)

import Data.Streaming.Text

-- | A specific character encoding.

--

-- Since 0.3.0

data Codec = Codec

    { _codecName :: T.Text

    , codecEncode

        :: T.Text

        -> (B.ByteString, Maybe (TextException, T.Text))

    , codecDecode

        :: B.ByteString

        -> (T.Text, Either

            (TextException, B.ByteString)

            B.ByteString)

    }

    | NewCodec T.Text (T.Text -> B.ByteString) (B.ByteString -> DecodeResult)

instance Show Codec where

    showsPrec d c =

        let (cnst, name) = case c of

                Codec t _ _    -> ("Codec ", t)

                NewCodec t _ _ -> ("NewCodec ", t)

        in showParen (d > 10) $ showString cnst . shows name

-- | Emit each line separately

--

-- Since 0.4.1

lines :: Monad m => Conduit T.Text m T.Text

lines =

    awaitText T.empty

  where

    awaitText buf = await >>= maybe (finish buf) (process buf)

    finish buf = unless (T.null buf) (yield buf)

    process buf text = yieldLines $ buf `T.append` text

    yieldLines buf =

      let (line, rest) = T.break (== '\n') buf

      in  case T.uncons rest of

            Just (_, rest') -> yield line >> yieldLines rest'

            _ -> awaitText line

-- | Variant of the lines function with an integer parameter.

-- The text length of any emitted line

-- never exceeds the value of the parameter. Whenever

-- this is about to happen a LengthExceeded exception

-- is thrown. This function should be used instead

-- of the lines function whenever we are dealing with

-- user input (e.g. a file upload) because we can't be sure that

-- user input won't have extraordinarily large lines which would

-- require large amounts of memory if consumed.

linesBounded :: MonadThrow m => Int -> Conduit T.Text m T.Text

linesBounded maxLineLen =

    awaitText 0 T.empty

  where

    awaitText len buf = await >>= maybe (finish buf) (process len buf)

    finish buf = unless (T.null buf) (yield buf)

    process len buf text =

      let (line, rest) = T.break (== '\n') text

          len' = len + T.length line

      in  if len' > maxLineLen

            then lift $ monadThrow (LengthExceeded maxLineLen)

            else case T.uncons rest of

                   Just (_, rest') ->

                     yield (buf `T.append` line) >> process 0 T.empty rest'

                   _ ->

                     awaitText len' $ buf `T.append` text

-- | Convert text into bytes, using the provided codec. If the codec is

-- not capable of representing an input character, an exception will be thrown.

--

-- Since 0.3.0

encode :: MonadThrow m => Codec -> Conduit T.Text m B.ByteString

encode (NewCodec _ enc _) = CL.map enc

encode codec = CL.mapM $ \t -> do

    let (bs, mexc) = codecEncode codec t

    maybe (return bs) (monadThrow . fst) mexc

decodeNew

    :: Monad m

    => (Int -> B.ByteString -> T.Text -> B.ByteString -> Conduit B.ByteString m T.Text)

    -> t

    -> Int

    -> (B.ByteString -> DecodeResult)

    -> Conduit B.ByteString m T.Text

decodeNew onFailure _name =

    loop

  where

    loop consumed dec =

        await >>= maybe finish go

      where

        finish =

            case dec B.empty of

                DecodeResultSuccess _ _ -> return ()

                DecodeResultFailure t rest -> onFailure consumed B.empty t rest

        {-# INLINE finish #-}

        go bs | B.null bs = loop consumed dec

        go bs =

            case dec bs of

                DecodeResultSuccess t dec' -> do

                    let consumed' = consumed + B.length bs

                        next = do

                            unless (T.null t) (yield t)

                            loop consumed' dec'

                     in consumed' `seq` next

                DecodeResultFailure t rest -> onFailure consumed bs t rest

-- | Decode a stream of UTF8 data, and replace invalid bytes with the Unicode

-- replacement character.

--

-- Since 1.1.1

decodeUtf8Lenient :: Monad m => Conduit B.ByteString m T.Text

decodeUtf8Lenient =

    decodeNew onFailure "UTF8-lenient" 0 Data.Streaming.Text.decodeUtf8

  where

    onFailure _consumed _bs t rest = do

        unless (T.null t) (yield t)

        case B.uncons rest of

            Nothing -> return ()

            Just (_, rest') -> do

                unless (B.null rest') (leftover rest')

                yield $ T.singleton '\xFFFD'

        decodeUtf8Lenient

-- | Convert bytes into text, using the provided codec. If the codec is

-- not capable of decoding an input byte sequence, an exception will be thrown.

--

-- Since 0.3.0

decode :: MonadThrow m => Codec -> Conduit B.ByteString m T.Text

decode (NewCodec name _ start) =

    decodeNew onFailure name 0 start

  where

    onFailure consumed bs t rest = do

        unless (T.null t) (yield t)

        leftover rest -- rest will never be null, no need to check

        let consumed' = consumed + B.length bs - B.length rest

        monadThrow $ NewDecodeException name consumed' (B.take 4 rest)

    {-# INLINE onFailure #-}

decode codec =

    loop id

  where

    loop front = await >>= maybe (finish front) (go front)

    finish front =

        case B.uncons $ front B.empty of

            Nothing -> return ()

            Just (w, _) -> lift $ monadThrow $ DecodeException codec w

    go front bs' =

        case extra of

            Left (exc, _) -> lift $ monadThrow exc

            Right bs'' -> yield text >> loop (B.append bs'')

      where

        (text, extra) = codecDecode codec bs

        bs = front bs'

-- |

-- Since 0.3.0

data TextException = DecodeException Codec Word8

                   | EncodeException Codec Char

                   | LengthExceeded Int

                   | TextException Exc.SomeException

                   | NewDecodeException !T.Text !Int !B.ByteString

    deriving Typeable

instance Show TextException where

    show (DecodeException codec w) = concat

        [ "Error decoding legacy Data.Conduit.Text codec "

        , show codec

        , " when parsing byte: "

        , show w

        ]

    show (EncodeException codec c) = concat

        [ "Error encoding legacy Data.Conduit.Text codec "

        , show codec

        , " when parsing char: "

        , show c

        ]

    show (LengthExceeded i) = "Data.Conduit.Text.linesBounded: line too long: " ++ show i

    show (TextException se) = "Data.Conduit.Text.TextException: " ++ show se

    show (NewDecodeException codec consumed next) = concat

        [ "Data.Conduit.Text.decode: Error decoding stream of "

        , T.unpack codec

        , " bytes. Error encountered in stream at offset "

        , show consumed

        , ". Encountered at byte sequence "

        , show next

        ]

instance Exc.Exception TextException

-- |

-- Since 0.3.0

utf8 :: Codec

utf8 = NewCodec (T.pack "UTF-8") TE.encodeUtf8 Data.Streaming.Text.decodeUtf8

-- |

-- Since 0.3.0

utf16_le :: Codec

utf16_le = NewCodec (T.pack "UTF-16-LE") TE.encodeUtf16LE decodeUtf16LE

-- |

-- Since 0.3.0

utf16_be :: Codec

utf16_be = NewCodec (T.pack "UTF-16-BE") TE.encodeUtf16BE decodeUtf16BE

-- |

-- Since 0.3.0

utf32_le :: Codec

utf32_le = NewCodec (T.pack "UTF-32-LE") TE.encodeUtf32LE decodeUtf32LE

-- |

-- Since 0.3.0

utf32_be :: Codec

utf32_be = NewCodec (T.pack "UTF-32-BE") TE.encodeUtf32BE decodeUtf32BE

-- |

-- Since 0.3.0

ascii :: Codec

ascii = Codec name enc dec where

    name = T.pack "ASCII"

    enc text = (bytes, extra) where

        (safe, unsafe) = T.span (\c -> ord c <= 0x7F) text

        bytes = B8.pack (T.unpack safe)

        extra = if T.null unsafe

            then Nothing

            else Just (EncodeException ascii (T.head unsafe), unsafe)

    dec bytes = (text, extra) where

        (safe, unsafe) = B.span (<= 0x7F) bytes

        text = T.pack (B8.unpack safe)

        extra = if B.null unsafe

            then Right B.empty

            else Left (DecodeException ascii (B.head unsafe), unsafe)

-- |

-- Since 0.3.0

iso8859_1 :: Codec

iso8859_1 = Codec name enc dec where

    name = T.pack "ISO-8859-1"

    enc text = (bytes, extra) where

        (safe, unsafe) = T.span (\c -> ord c <= 0xFF) text

        bytes = B8.pack (T.unpack safe)

        extra = if T.null unsafe

            then Nothing

            else Just (EncodeException iso8859_1 (T.head unsafe), unsafe)

    dec bytes = (T.pack (B8.unpack bytes), Right B.empty)

-- |

--

-- Since 1.0.8

takeWhile :: Monad m

          => (Char -> Bool)

          -> Conduit T.Text m T.Text

takeWhile p =

    loop

  where

    loop = await >>= maybe (return ()) go

    go t =

        case T.span p t of

            (x, y)

                | T.null y -> yield x >> loop

                | otherwise -> yield x >> leftover y

-- |

--

-- Since 1.0.8

dropWhile :: Monad m

          => (Char -> Bool)

          -> Consumer T.Text m ()

dropWhile p =

    loop

  where

    loop = await >>= maybe (return ()) go

    go t

        | T.null x = loop

        | otherwise = leftover x

      where

        x = T.dropWhile p t

-- |

--

-- Since 1.0.8

take :: Monad m => Int -> Conduit T.Text m T.Text

take =

    loop

  where

    loop i = await >>= maybe (return ()) (go i)

    go i t

        | diff == 0 = yield t

        | diff < 0 =

            let (x, y) = T.splitAt i t

             in yield x >> leftover y

        | otherwise = yield t >> loop diff

      where

        diff = i - T.length t

-- |

--

-- Since 1.0.8

drop :: Monad m => Int -> Consumer T.Text m ()

drop =

    loop

  where

    loop i = await >>= maybe (return ()) (go i)

    go i t

        | diff == 0 = return ()

        | diff < 0 = leftover $ T.drop i t

        | otherwise = loop diff

      where

        diff = i - T.length t

-- |

--

-- Since 1.0.8

foldLines :: Monad m

          => (a -> ConduitM T.Text o m a)

          -> a

          -> ConduitM T.Text o m a

foldLines f =

    start

  where

    start a = CL.peek >>= maybe (return a) (const $ loop $ f a)

    loop consumer = do

        a <- takeWhile (/= '\n') =$= do

            a <- CL.map (T.filter (/= '\r')) =$= consumer

            CL.sinkNull

            return a

        drop 1

        start a

-- |

--

-- Since 1.0.8

withLine :: Monad m

         => Sink T.Text m a

         -> Consumer T.Text m (Maybe a)

withLine consumer = toConsumer $ do

    mx <- CL.peek

    case mx of

        Nothing -> return Nothing

        Just _ -> do

            x <- takeWhile (/= '\n') =$ do

                x <- CL.map (T.filter (/= '\r')) =$ consumer

                CL.sinkNull

                return x

            drop 1

            return $ Just x

-- | Decode a stream of UTF8-encoded bytes into a stream of text, throwing an

-- exception on invalid input.

--

-- Since 1.0.15

decodeUtf8 :: MonadThrow m => Conduit B.ByteString m T.Text

decodeUtf8 = decode utf8

    {- no meaningful performance advantage

    CI.ConduitM (loop 0 decodeUtf8)

  where

    loop consumed dec =

        CI.NeedInput go finish

      where

        finish () =

            case dec B.empty of

                DecodeResultSuccess _ _ -> return ()

                DecodeResultFailure t rest -> onFailure B.empty t rest

        {-# INLINE finish #-}

        go bs | B.null bs = CI.NeedInput go finish

        go bs =

            case dec bs of

                DecodeResultSuccess t dec' -> do

                    let consumed' = consumed + B.length bs

                        next' = loop consumed' dec'

                        next

                            | T.null t = next'

                            | otherwise = CI.HaveOutput next' (return ()) t

                     in consumed' `seq` next

                DecodeResultFailure t rest -> onFailure bs t rest

        onFailure bs t rest = do

            unless (T.null t) (CI.yield t)

            unless (B.null rest) (CI.leftover rest)

            let consumed' = consumed + B.length bs - B.length rest

            monadThrow $ NewDecodeException (T.pack "UTF-8") consumed' (B.take 4 rest)

        {-# INLINE onFailure #-}

    -}

{-# INLINE decodeUtf8 #-}

-- | Encode a stream of text into a stream of bytes.

--

-- Since 1.0.15

encodeUtf8 :: Monad m => Conduit T.Text m B.ByteString

encodeUtf8 = CL.map TE.encodeUtf8

{-# INLINE encodeUtf8 #-}

-- | Automatically determine which UTF variant is being used. This function

-- checks for BOMs, removing them as necessary. It defaults to assuming UTF-8.

--

-- Since 1.1.9

detectUtf :: MonadThrow m => Conduit B.ByteString m T.Text

detectUtf =

    go id

  where

    go front = await >>= maybe (close front) (push front)

    push front bs'

        | B.length bs < 4 = go $ B.append bs

        | otherwise       = leftDecode bs

      where bs = front bs'

    close front = leftDecode $ front B.empty

    leftDecode bs = leftover bsOut >> decode codec

      where

        bsOut = B.append (B.drop toDrop x) y

        (x, y) = B.splitAt 4 bs

        (toDrop, codec) =

            case B.unpack x of

                [0x00, 0x00, 0xFE, 0xFF] -> (4, utf32_be)

                [0xFF, 0xFE, 0x00, 0x00] -> (4, utf32_le)

                0xFE : 0xFF: _           -> (2, utf16_be)

                0xFF : 0xFE: _           -> (2, utf16_le)

                0xEF : 0xBB: 0xBF : _    -> (3, utf8)

                _                        -> (0, utf8) -- Assuming UTF-8

{-# INLINE detectUtf #-}