{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE OverloadedStrings #-}

-- | Utility functions for running a parser against a socket

module Network.Attoparsec (parseMany, parseOne) where

import           Control.Monad.Error

import qualified Data.ByteString            as BS
import qualified Network.Socket             as NS
import qualified Network.Socket.ByteString  as NSB
import qualified Data.Attoparsec.ByteString as Atto

-- | The parsing continuation form of a "Data.Attoparsec" parser.
type ParseC a = BS.ByteString -> Atto.Result a

-- | The type of parsing to perform, greedy or non-greedy
data ParseMode = Single | Many

-- | Consumes input from socket and attempts to parse a structure. This function
--   will terminate if one of three conditions is met:
--
--   * the parser has completed succesfully;
--   * the parser failed with invalid data;
--   * the connection is closed.
--
--   Depending upon the available data on the socket, this function might block.
parseBuffer :: ( MonadIO m
               , MonadError String m
               , Show a)
            => ParseC a          -- ^ Initial parser state
            -> ParseMode         -- ^ Whether to perform greedy or non-greedy parsing
            -> BS.ByteString     -- ^ Unconsumed buffer from previous run
            -> ParseC a          -- ^ Current parser state
            -> m (ParseC a, [a]) -- ^ Next parser state with parsed values
parseBuffer p0 mode =

  let next bCur pCur =
        case pCur bCur of
         -- On error, throw error through MonadError
         Atto.Fail    err _ _  -> throwError ("An error occured while parsing input: " ++ show err)
         Atto.Partial p1       -> return (p1, [])
         Atto.Done    b1 v     ->
           if BS.null b1
             -- This means a "perfect parse" occured: exactly enough data was on
             -- the socket to complete one parse round.
             then return (p0, [v])
             else case mode of
                   -- We are in single-object parsing mode, have parsed one object,
                   -- but still have data left on the buffer: at this point, we can
                   -- either discard the data on the buffer, or throw an error.
                   --
                   -- We throw an error, since within "single-object parsing mode"
                   -- we assume only perfect parses happen.
                   Single -> throwError ("Unconsumed data left on socket: " ++ show b1)

                   -- Multi-object parsing mode, in which case we will enter
                   -- recursion.
                   Many   -> do
                     (p1, xs) <- next b1 p0
                     return (p1, v : xs)

  in next

-- | Incrementally reads data from socket and parses as many objects as possible
parseMany :: ( MonadIO m
             , MonadError String m
             , Show a)
          => NS.Socket         -- ^ Socket to read data from
          -> ParseC a          -- ^ Initial parser state
          -> ParseC a          -- ^ Continuation parser state
          -> m (ParseC a, [a]) -- ^ Next parser state with parsed values
parseMany s p0 pCur = do
  buf <- liftIO $ NSB.recv s 4096
  (p1, xs) <- parseBuffer p0 Many buf pCur
  return (p1, xs)

-- | Similar to parseMany, but assumes that there will only be enough data for a
--   single succesful parse on the socket, and guarantees that exactly one item
--   will be parsed.
--
--   __Warning:__ this function will /not/ work correctly when input data is
--   pipelined. The parser might consume more data than required from the socket,
--   or a partial second object is parsed, and the parser state and buffer will
--   be discarded.
parseOne :: ( MonadIO m
            , MonadError String m
            , Show a)
         => NS.Socket -- ^ Socket to read data from
         -> ParseC a  -- ^ Initial parser state
         -> m a       -- ^ Parsed value
parseOne s p0 = do
  buf <- liftIO $ NSB.recv s 4096
  (p1, value) <- parseBuffer p0 Single buf p0

  case value of
   -- We do not yet have enough data for a single item, let's request more
   []  -> parseOne s p1
   [x] -> return x

   -- This is an internal error, since it means our single-object parser
   -- returned multiple objects.
   _   -> error ("More than one element parsed: " ++ show value)