{-# LANGUAGE OverloadedStrings, CPP #-}

module Network.Wai.Application.Classic.EventSource (
    bodyToEventSource
  ) where

import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.ByteString.Builder
import Data.ByteString.Char8 ()
import Data.Conduit
import qualified Data.Conduit.List as CL
import qualified Network.HTTP.Client as H
import qualified Network.HTTP.Client.Conduit as HC

lineBreak :: ByteString -> Int -> Maybe Int
lineBreak :: ByteString -> Int -> Maybe Int
lineBreak ByteString
bs Int
n = Maybe Int
go
  where
    len :: Int
len = ByteString -> Int
BS.length ByteString
bs
    go :: Maybe Int
go | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Maybe Int
forall a. Maybe a
Nothing
       | Bool
otherwise = case ByteString
bs ByteString -> Int -> Word8
`BS.index` Int
n of
                         Word8
13 -> Int -> Maybe Int
go' (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                         Word8
10 -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                         Word8
_ -> Maybe Int
forall a. Maybe a
Nothing
    go' :: Int -> Maybe Int
go' Int
n' | Int
n' Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
len = Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n'
           | Bool
otherwise = case ByteString
bs ByteString -> Int -> Word8
`BS.index` Int
n' of
                             Word8
10 -> Int -> Maybe Int
forall a. a -> Maybe a
Just (Int
n'Int -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                             Word8
_ -> Int -> Maybe Int
forall a. a -> Maybe a
Just Int
n'

-- splitDoubleLineBreak "aaa\n\nbbb" == ["aaa\n\n", "bbb"]
-- splitDoubleLineBreak "aaa\n\nbbb\n\n" == ["aaa\n\n", "bbb\n\n", ""]
-- splitDoubleLineBreak "aaa\r\n\rbbb\n\r\n" == ["aaa\r\n\r", "bbb\n\r\n", ""]
-- splitDoubleLineBreak "aaa" == ["aaa"]
-- splitDoubleLineBreak "" == [""]
splitDoubleLineBreak :: ByteString -> [ByteString]
splitDoubleLineBreak :: ByteString -> [ByteString]
splitDoubleLineBreak ByteString
str = ByteString -> Int -> [ByteString]
go ByteString
str Int
0
  where
    go :: ByteString -> Int -> [ByteString]
go ByteString
bs Int
n | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< ByteString -> Int
BS.length ByteString
str =
                    case ByteString -> Int -> Maybe Int
lineBreak ByteString
bs Int
n of
                        Maybe Int
Nothing -> ByteString -> Int -> [ByteString]
go ByteString
bs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                        Just Int
n' ->
                            case ByteString -> Int -> Maybe Int
lineBreak ByteString
bs Int
n' of
                                Maybe Int
Nothing -> ByteString -> Int -> [ByteString]
go ByteString
bs (Int
nInt -> Int -> Int
forall a. Num a => a -> a -> a
+Int
1)
                                Just Int
n'' ->
                                    let (ByteString
xs,ByteString
ys) = Int -> ByteString -> (ByteString, ByteString)
BS.splitAt Int
n'' ByteString
bs
                                    in ByteString
xsByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
:ByteString -> Int -> [ByteString]
go ByteString
ys Int
0
            | Bool
otherwise = [ByteString
bs]

#if MIN_VERSION_conduit(1,3,0)
eventSourceConduit :: ConduitT ByteString (Flush Builder) IO ()
#else
eventSourceConduit :: Conduit ByteString IO (Flush Builder)
#endif
eventSourceConduit :: ConduitT ByteString (Flush Builder) IO ()
eventSourceConduit = (ByteString -> ByteString -> (ByteString, [Flush Builder]))
-> ByteString -> ConduitT ByteString (Flush Builder) IO ()
forall (m :: * -> *) a accum b.
Monad m =>
(a -> accum -> (accum, [b])) -> accum -> ConduitT a b m ()
CL.concatMapAccum ByteString -> ByteString -> (ByteString, [Flush Builder])
f ByteString
""
  where
    f :: ByteString -> ByteString -> (ByteString, [Flush Builder])
f ByteString
input ByteString
rest = ([ByteString] -> ByteString
forall a. [a] -> a
last [ByteString]
xs, (ByteString -> [Flush Builder]) -> [ByteString] -> [Flush Builder]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap ByteString -> [Flush Builder]
addFlush ([ByteString] -> [Flush Builder])
-> [ByteString] -> [Flush Builder]
forall a b. (a -> b) -> a -> b
$ [ByteString] -> [ByteString]
forall a. [a] -> [a]
init [ByteString]
xs)
      where
        addFlush :: ByteString -> [Flush Builder]
addFlush ByteString
x = [Builder -> Flush Builder
forall a. a -> Flush a
Chunk (ByteString -> Builder
byteString ByteString
x), Flush Builder
forall a. Flush a
Flush]
        xs :: [ByteString]
xs = ByteString -> [ByteString]
splitDoubleLineBreak (ByteString
rest ByteString -> ByteString -> ByteString
`BS.append` ByteString
input)

-- insert Flush if exists a double line-break
#if MIN_VERSION_conduit(1,3,0)
bodyToEventSource :: H.BodyReader -> ConduitT () (Flush Builder) IO ()
bodyToEventSource :: BodyReader -> ConduitT () (Flush Builder) IO ()
bodyToEventSource BodyReader
br = BodyReader -> ConduitM () ByteString IO ()
forall (m :: * -> *) i.
MonadIO m =>
BodyReader -> ConduitM i ByteString m ()
HC.bodyReaderSource BodyReader
br ConduitM () ByteString IO ()
-> ConduitT ByteString (Flush Builder) IO ()
-> ConduitT () (Flush Builder) IO ()
forall (m :: * -> *) a b c r.
Monad m =>
ConduitM a b m () -> ConduitM b c m r -> ConduitM a c m r
.| ConduitT ByteString (Flush Builder) IO ()
eventSourceConduit
#else
bodyToEventSource :: H.BodyReader -> Source IO (Flush Builder)
bodyToEventSource br = HC.bodyReaderSource br $= eventSourceConduit
#endif