-- |
-- Module:     Network.IHttp.Header
-- Copyright:  (c) 2010 Ertugrul Soeylemez
-- License:    BSD3
-- Maintainer: Ertugrul Soeylemez <es@ertes.de>
-- Stability:  beta
--
-- Iteratees for header parsing.

{-# LANGUAGE OverloadedStrings, ScopedTypeVariables #-}

module Network.IHttp.Header
    ( -- * Iteratees
      httpHeader,
      httpHeaders,

      -- * Enumerators
      enumHeaders
    )
    where

import qualified Data.ByteString as B
import qualified Data.Map as M
import Control.Applicative
import Data.ByteString (ByteString)
import Data.ByteString.Char8 ()
import Data.Enumerator as E
import Data.Enumerator.List as EL
import Network.IHttp.Parsers
import Network.IHttp.Tools
import Network.IHttp.Types


-- | Enumerate a 'HeaderMap' as a protocol string stream.  You can use
-- 'Data.Enumerator.Binary.iterHandle' to send it.  Note that this
-- enumerator never generates continuation lines.  It also does not
-- enumerate the final empty line.

enumHeaders :: forall b m. Monad m => HeaderMap -> Enumerator ByteString m b
enumHeaders headers = enum
    where
    format :: ByteString -> ByteString -> [ByteString] -> [ByteString]
    format header content rest = header : ": " : content : "\r\n" : rest

    enum :: Enumerator ByteString m b
    enum (Continue k) = k (Chunks . M.foldrWithKey format [] $ headers)
    enum step         = returnI step


-- | Get the next header from the 'netLinesEmpty'-splitted stream.  The
-- header's content is length-limited by the given argument.  If it's
-- longer, it's safely truncated in constant space.  This iteratee
-- throws an iteratee error, if the next lines are not a valid HTTP
-- header or the stream ends prematurely.  If the next line is an empty
-- line, this iteratee returns 'Nothing'.

httpHeader ::
    forall m. Monad m =>
    Int -> Iteratee ByteString m (Maybe (ByteString, ByteString))
httpHeader n = do
    line <- EL.head >>=
            maybe (throwError $ InvalidHeaderError "Premature end of stream") return
    if B.null line
      then return Nothing
      else do
          (hdrName, hdrTxt) <- parseIter httpFirstHeaderP InvalidHeaderError line
          Just <$> continue (loop hdrName (B.take n hdrTxt))

    where
    loop :: ByteString -> ByteString -> Stream ByteString ->
            Iteratee ByteString m (ByteString, ByteString)
    loop hdrName hdrTxt EOF = yield (hdrName, hdrTxt) EOF
    loop hdrName hdrTxt (Chunks []) = continue (loop hdrName hdrTxt)
    loop hdrName hdrTxt' ch@(Chunks (line:lines)) =
        let (pfx, sfx) = B.span (\c -> c == 32 || c == 9) line
            hdrTxt = B.take n $ B.append (B.snoc hdrTxt' 32) sfx
        in if B.null pfx
             then yield (hdrName, hdrTxt') ch
             else hdrTxt `seq` loop hdrName hdrTxt (Chunks lines)


-- | Get the headers of an HTTP request from a 'netLinesEmpty'-splitted
-- byte stream.  The first 'Int' specifies the maximum length of
-- individual headers.  The second 'Int' specifies the maximum number of
-- headers.  This iteratee throws an iteratee error on invalid input, of
-- if the stream ends prematurely.
--
-- Excess data is truncated safely in constant space.

httpHeaders :: forall m. Monad m => Int -> Int -> Iteratee ByteString m HeaderMap
httpHeaders maxLine maxHeaders =
    loop M.empty

    where
    loop :: HeaderMap -> Iteratee ByteString m HeaderMap
    loop m' = do
        mHeader <- httpHeader maxLine
        case mHeader of
          Just (hdrName, hdrSfx) -> do
              let accum hdrPfx = B.take maxLine $ B.concat [hdrPfx, ", ", hdrSfx]
                  hdrTxt = maybe hdrSfx accum (M.lookup hdrName m')
              if M.size m' < maxHeaders
                then hdrTxt `seq` loop (M.insert hdrName hdrTxt m')
                else loop m'
          Nothing -> return m'