---------------------------------------------------------
-- |
-- Module        : Data.ByteString.Lazy.Util
-- Copyright     : Michael Snoyman
-- License       : BSD3
--
-- Maintainer    : Michael Snoyman <michael@snoyman.com>
-- Stability     : Unstable
-- Portability   : portable
--
-- Various utilities to assist in dealing with lazy bytestrings.
--
---------------------------------------------------------

module Data.ByteString.Lazy.Util
    ( ord
    , stripPrefix
    , breakAt
    , breakAtString
    , takeLine
    , chompBS
    , takeUntilBlank
    ) where

import qualified Data.ByteString.Lazy as BS
import qualified Data.Char as C
import Data.Word (Word8)

-- | Get the ASCII value of character. Differers from regular ord in that
-- it returns an Integral, so it is automatically cast to eg a Word8.
ord :: Integral a => Char -> a
ord = fromInteger . toInteger . C.ord

-- | Strip a prefix from a bytestring if it's there.
stripPrefix :: Word8 -> BS.ByteString -> BS.ByteString
stripPrefix p bs
    | BS.null bs = bs
    | BS.head bs == p = BS.tail bs
    | otherwise = bs

-- | Break a bytestring into two at the first occurence of the given 'Word8'.
-- That 'Word8' should not appear in either piece.
breakAt :: Word8 -> BS.ByteString -> (BS.ByteString, BS.ByteString)
breakAt p bs =
    let (x, y) = BS.span (/= p) bs
        y' = stripPrefix p y
     in (x, y')

-- | Same as 'breakAt', but use a bytestring instead of a 'Word8'.
breakAtString :: BS.ByteString
              -> BS.ByteString
              -> (BS.ByteString, BS.ByteString)
breakAtString p c
    | BS.null c = (BS.empty, BS.empty)
    | p `BS.isPrefixOf` c = (BS.empty, BS.drop (BS.length p) c)
    | otherwise =
        let x = BS.head c
            xs = BS.tail c
            (next, rest) = breakAtString p xs
        in (BS.cons' x next, rest)

-- | Take a single line from a bytestring.
takeLine :: BS.ByteString -> (BS.ByteString, BS.ByteString)
takeLine bs =
    let (x, y) = BS.span (/= ord '\n') bs
        x' = if not (BS.null x) && BS.last x == ord '\r' then BS.init x else x
        y' = if not (BS.null y) && BS.head y == ord '\n' then BS.tail y else y
     in (x', y')

-- | Removes newline characters from the end of a string.
chompBS :: BS.ByteString -> BS.ByteString
chompBS s
    | BS.null s = s
    | BS.last s == ord '\n' =
        if BS.length s == 1 || BS.last (BS.init s) /= ord '\r'
            then BS.init s
            else BS.init (BS.init s)
    | BS.last s == ord '\r' = BS.init s
    | otherwise = s

-- | Take each line until the first blank line and return as first.
-- The rest of the content is returned as second.
takeUntilBlank :: BS.ByteString -> ([BS.ByteString], BS.ByteString)
takeUntilBlank bs =
    let (next, rest) = takeLine bs
     in if BS.null next
            then ([], rest)
            else let (nexts, rest') = takeUntilBlank rest
                  in (next : nexts, rest')