{-# LANGUAGE Rank2Types #-}

{-# LANGUAGE BangPatterns #-}

{-# LANGUAGE OverloadedStrings #-}

-- |

-- Module      : Network.Socks5.Parse

-- License     : BSD-style

-- Maintainer  : Vincent Hanquez <vincent@snarc.org>

-- Stability   : experimental

-- Portability : portable

--

-- A very simple bytestring parser related to Parsec and Attoparsec

--

-- Simple example:

--

-- > > parse ((,) <$> take 2 <*> byte 0x20 <*> (bytes "abc" *> anyByte)) "xx abctest"

-- > ParseOK "est" ("xx", 116)

--

module Network.Socks5.Parse

    ( Parser

    , Result(..)

    -- * run the Parser

    , parse

    , parseFeed

    -- * Parser methods

    , byte

    , anyByte

    , bytes

    , take

    , takeWhile

    , takeAll

    , skip

    , skipWhile

    , skipAll

    , takeStorable

    ) where

import Control.Applicative

import Control.Monad

import Data.ByteString (ByteString)

import qualified Data.ByteString as B

import qualified Data.ByteString.Internal as B (toForeignPtr)

import Data.Word

import Foreign.Storable (Storable, peekByteOff, sizeOf)

import Foreign.ForeignPtr (withForeignPtr)

import Prelude hiding (take, takeWhile)

import System.IO.Unsafe (unsafePerformIO)

-- | Simple parsing result, that represent respectively:

--

-- * failure: with the error message

--

-- * continuation: that need for more input data

--

-- * success: the remaining unparsed data and the parser value

data Result a =

      ParseFail String

    | ParseMore (ByteString -> Result a)

    | ParseOK   ByteString a

instance Show a => Show (Result a) where

    show (ParseFail err) = "ParseFailure: " ++ err

    show (ParseMore _)   = "ParseMore _"

    show (ParseOK b a)   = "ParseOK " ++ show a ++ " " ++ show b

type Failure r = ByteString -> String -> Result r

type Success a r = ByteString -> a -> Result r

-- | Simple ByteString parser structure

newtype Parser a = Parser

    { runParser :: forall r . ByteString -> Failure r -> Success a r -> Result r }

instance Monad Parser where

    fail errorMsg = Parser $ \buf err _ -> err buf ("failed: " ++ errorMsg)

    return v = Parser $ \buf _ ok -> ok buf v

    m >>= k = Parser $ \buf err ok ->

         runParser m buf err (\buf' a -> runParser (k a) buf' err ok)

instance MonadPlus Parser where

    mzero = fail "Parser.MonadPlus.mzero"

    mplus f g = Parser $ \buf err ok ->

        -- rewrite the err callback of @f to call @g

        runParser f buf (\_ _ -> runParser g buf err ok) ok

instance Functor Parser where

    fmap f p = Parser $ \buf err ok ->

        runParser p buf err (\b a -> ok b (f a))

instance Applicative Parser where

    pure      = return

    (<*>) d e = d >>= \b -> e >>= \a -> return (b a)

instance Alternative Parser where

    empty = fail "Parser.Alternative.empty"

    (<|>) = mplus

-- | Run a parser on an @initial ByteString.

--

-- If the Parser need more data than available, the @feeder function

-- is automatically called and fed to the More continuation.

parseFeed :: Monad m => m B.ByteString -> Parser a -> B.ByteString -> m (Result a)

parseFeed feeder p initial = loop $ parse p initial

  where loop (ParseMore k) = feeder >>= (loop . k)

        loop r             = return r

-- | Run a Parser on a ByteString and return a 'Result'

parse :: Parser a -> ByteString -> Result a

parse p s = runParser p s (\_ msg -> ParseFail msg) (\b a -> ParseOK b a)

------------------------------------------------------------

getMore :: Parser ()

getMore = Parser $ \buf err ok -> ParseMore $ \nextChunk ->

    if B.null nextChunk

        then err buf "EOL: need more data"

        else ok (B.append buf nextChunk) ()

getAll :: Parser ()

getAll = Parser $ \buf err ok -> ParseMore $ \nextChunk ->

    if B.null nextChunk

        then ok buf ()

        else runParser getAll (B.append buf nextChunk) err ok

flushAll :: Parser ()

flushAll = Parser $ \buf err ok -> ParseMore $ \nextChunk ->

    if B.null nextChunk

        then ok buf ()

        else runParser getAll B.empty err ok

------------------------------------------------------------

-- | Get the next byte from the parser

anyByte :: Parser Word8

anyByte = Parser $ \buf err ok ->

    case B.uncons buf of

        Nothing      -> runParser (getMore >> anyByte) buf err ok

        Just (c1,b2) -> ok b2 c1

-- | Parse a specific byte at current position

--

-- if the byte is different than the expected on,

-- this parser will raise a failure.

byte :: Word8 -> Parser ()

byte w = Parser $ \buf err ok ->

    case B.uncons buf of

        Nothing      -> runParser (getMore >> byte w) buf err ok

        Just (c1,b2) | c1 == w   -> ok b2 ()

                     | otherwise -> err buf ("byte " ++ show w ++ " : failed")

-- | Parse a sequence of bytes from current position

--

-- if the following bytes don't match the expected

-- bytestring completely, the parser will raise a failure

bytes :: ByteString -> Parser ()

bytes allExpected = consumeEq allExpected

  where errMsg = "bytes " ++ show allExpected ++ " : failed"

        -- partially consume as much as possible or raise an error.

        consumeEq expected = Parser $ \actual err ok ->

            let eLen = B.length expected in

            if B.length actual >= eLen

                then    -- enough data for doing a full match

                        let (aMatch,aRem) = B.splitAt eLen actual

                         in if aMatch == expected

                                then ok aRem ()

                                else err actual errMsg

                else    -- not enough data, match as much as we have, and then recurse.

                        let (eMatch, eRem) = B.splitAt (B.length actual) expected

                         in if actual == eMatch

                                then runParser (getMore >> consumeEq eRem) B.empty err ok

                                else err actual errMsg

------------------------------------------------------------

-- | Take a storable from the current position in the stream

takeStorable :: Storable d

             => Parser d

takeStorable = anyStorable undefined

  where

    anyStorable :: Storable d => d -> Parser d

    anyStorable a = do

        (fptr, off, _) <- B.toForeignPtr <$> take (sizeOf a)

        return $ unsafePerformIO $ withForeignPtr fptr $ \ptr -> peekByteOff ptr off

-- | Take @n bytes from the current position in the stream

take :: Int -> Parser ByteString

take n = Parser $ \buf err ok ->

    if B.length buf >= n

        then let (b1,b2) = B.splitAt n buf in ok b2 b1

        else runParser (getMore >> take n) buf err ok

-- | Take bytes while the @predicate hold from the current position in the stream

takeWhile :: (Word8 -> Bool) -> Parser ByteString

takeWhile predicate = Parser $ \buf err ok ->

    case B.span predicate buf of

        (_, b2) | B.null b2 -> runParser (getMore >> takeWhile predicate) buf err ok

        (b1, b2) -> ok b2 b1

-- | Take the remaining bytes from the current position in the stream

takeAll :: Parser ByteString

takeAll = Parser $ \buf err ok ->

    runParser (getAll >> returnBuffer) buf err ok

  where

    returnBuffer = Parser $ \buf _ ok -> ok B.empty buf

-- | Skip @n bytes from the current position in the stream

skip :: Int -> Parser ()

skip n = Parser $ \buf err ok ->

    if B.length buf >= n

        then ok (B.drop n buf) ()

        else runParser (getMore >> skip (n - B.length buf)) B.empty err ok

-- | Skip bytes while the @predicate hold from the current position in the stream

skipWhile :: (Word8 -> Bool) -> Parser ()

skipWhile p = Parser $ \buf err ok ->

    case B.span p buf of

        (_, b2) | B.null b2 -> runParser (getMore >> skipWhile p) B.empty err ok

        (_, b2) -> ok b2 ()

-- | Skip all the remaining bytes from the current position in the stream

skipAll :: Parser ()

skipAll = Parser $ \buf err ok -> runParser flushAll buf err ok