content: {-# LANGUAGE RecordWildCards #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE MultiWayIf #-}
  module Network.HTTP.Client.Body
      ( makeChunkedReader
      , makeLengthReader
      , makeGzipReader
      , makeUnlimitedReader
      , brConsume
      , brEmpty
      , constBodyReader
      , brReadSome
      , brRead
      ) where
  
  import Network.HTTP.Client.Connection
  import Network.HTTP.Client.Types
  import Control.Exception (assert)
  import Data.ByteString (empty, uncons)
  import Data.IORef
  import qualified Data.ByteString as S
  import qualified Data.ByteString.Lazy as L
  import Control.Monad (unless, when)
  import qualified Data.Streaming.Zlib as Z
  
  -- | Get a single chunk of data from the response body, or an empty
  -- bytestring if no more data is available.
  --
  -- Note that in order to consume the entire request body, you will need to
  -- repeatedly call this function until you receive an empty @ByteString@ as a
  -- result.
  --
  -- Since 0.1.0
  brRead :: BodyReader -> IO S.ByteString
  brRead = id
  
  -- | Continuously call 'brRead', building up a lazy ByteString until a chunk is
  -- constructed that is at least as many bytes as requested.
  --
  -- Since 0.4.20
  brReadSome :: BodyReader -> Int -> IO L.ByteString
  brReadSome brRead' =
      loop id
    where
      loop front rem'
          | rem' <= 0 = return $ L.fromChunks $ front []
          | otherwise = do
              bs <- brRead'
              if S.null bs
                  then return $ L.fromChunks $ front []
                  else loop (front . (bs:)) (rem' - S.length bs)
  
  brEmpty :: BodyReader
  brEmpty = return S.empty
  
  constBodyReader :: [S.ByteString] -> IO BodyReader
  constBodyReader input = do
    iinput <- newIORef input
    return $ atomicModifyIORef iinput $ \input' ->
          case input' of
              [] -> ([], S.empty)
              x:xs -> (xs, x)
  
  -- | Strictly consume all remaining chunks of data from the stream.
  --
  -- Since 0.1.0
  brConsume :: BodyReader -> IO [S.ByteString]
  brConsume brRead' =
      go id
    where
      go front = do
          x <- brRead'
          if S.null x
              then return $ front []
              else go (front . (x:))
  
  makeGzipReader :: BodyReader -> IO BodyReader
  makeGzipReader brRead' = do
      inf <- Z.initInflate $ Z.WindowBits 31
      istate <- newIORef Nothing
      let goPopper popper = do
              res <- popper
              case res of
                  Z.PRNext bs -> do
                      writeIORef istate $ Just popper
                      return bs
                  Z.PRDone -> do
                      bs <- Z.flushInflate inf
                      if S.null bs
                          then start
                          else do
                              writeIORef istate Nothing
                              return bs
                  Z.PRError e -> throwHttp $ HttpZlibException e
          start = do
              bs <- brRead'
              if S.null bs
                  then return S.empty
                  else do
                      popper <- Z.feedInflate inf bs
                      goPopper popper
      return $ do
          state <- readIORef istate
          case state of
              Nothing -> start
              Just popper -> goPopper popper
  
  makeUnlimitedReader
    :: IO () -- ^ cleanup
    -> Connection
    -> IO BodyReader
  makeUnlimitedReader cleanup Connection {..} = do
      icomplete <- newIORef False
      return $ do
          bs <- connectionRead
          when (S.null bs) $ do
            writeIORef icomplete True
            cleanup
          return bs
  
  makeLengthReader
    :: IO () -- ^ cleanup
    -> Int
    -> Connection
    -> IO BodyReader
  makeLengthReader cleanup count0 Connection {..} = do
      icount <- newIORef count0
      return $ do
          count <- readIORef icount
          if count <= 0
              then return empty
              else do
                  bs <- connectionRead
                  when (S.null bs) $ throwHttp $ ResponseBodyTooShort (fromIntegral count0) (fromIntegral $ count0 - count)
                  case compare count $ S.length bs of
                      LT -> do
                          let (x, y) = S.splitAt count bs
                          connectionUnread y
                          writeIORef icount (-1)
                          cleanup
                          return x
                      EQ -> do
                          writeIORef icount (-1)
                          cleanup
                          return bs
                      GT -> do
                          writeIORef icount (count - S.length bs)
                          return bs
  
  makeChunkedReader
    :: Maybe MaxHeaderLength
    -> IO () -- ^ cleanup
    -> Bool -- ^ raw
    -> Connection
    -> IO BodyReader
  makeChunkedReader mhl cleanup raw conn@Connection {..} = do
      icount <- newIORef 0
      return $ do
        bs <- go icount
        when (S.null bs) cleanup
        pure bs
    where
      go icount = do
          count0 <- readIORef icount
          (rawCount, count) <-
              if count0 == 0
                  then readHeader
                  else return (empty, count0)
          if count <= 0
              then do
                  -- count == -1 indicates that all chunks have been consumed
                  writeIORef icount (-1)
                  if | count /= -1 && raw -> S.append rawCount <$> readTrailersRaw
                     | count /= -1        -> consumeTrailers *> pure empty
                     | otherwise          -> pure empty
              else do
                  (bs, count') <- readChunk count
                  writeIORef icount count'
                  return $ appendHeader rawCount bs
  
      appendHeader
        | raw = S.append
        | otherwise = flip const
  
      readChunk 0 = return (empty, 0)
      readChunk remainder = do
          bs <- connectionRead
          when (S.null bs) $ throwHttp InvalidChunkHeaders
          case compare remainder $ S.length bs of
              LT -> do
                  let (x, y) = S.splitAt remainder bs
                  assert (not $ S.null y) $ connectionUnread y
                  requireNewline
                  done x
              EQ -> do
                  requireNewline
                  done bs
              GT -> return (bs, remainder - S.length bs)
        where
          done x
            | raw = return (x `S.append` "\r\n", 0)
            | otherwise = return (x, 0)
  
      requireNewline = do
          bs <- connectionReadLine mhl conn
          unless (S.null bs) $ throwHttp InvalidChunkHeaders
  
      readHeader = do
          bs <- connectionReadLine mhl conn
          case parseHex bs of
              Nothing -> throwHttp InvalidChunkHeaders
              Just hex -> return (bs `S.append` "\r\n", hex)
  
      parseHex bs0 =
          case uncons bs0 of
              Just (w0, bs')
                  | Just i0 <- toI w0 -> Just $ parseHex' i0 bs'
              _ -> Nothing
      parseHex' i bs =
          case uncons bs of
              Just (w, bs')
                  | Just i' <- toI w -> parseHex' (i * 16 + i') bs'
              _ -> i
  
      toI w
          | 48 <= w && w <= 57  = Just $ fromIntegral w - 48
          | 65 <= w && w <= 70  = Just $ fromIntegral w - 55
          | 97 <= w && w <= 102 = Just $ fromIntegral w - 87
          | otherwise = Nothing
  
      readTrailersRaw = do
          bs <- connectionReadLine mhl conn
          if S.null bs
          then pure "\r\n"
          else (bs `S.append` "\r\n" `S.append`) <$> readTrailersRaw
  
      consumeTrailers = connectionDropTillBlankLine mhl conn
  