-----------------------------------------------------------------------------
-- |
-- Module: Data.Enumerator.List
-- Copyright: 2010 John Millikin
-- License: MIT
--
-- Maintainer: jmillikin@gmail.com
-- Portability: portable
--
-- This module is intended to be imported qualified:
--
-- @
-- import qualified Data.Enumerator.List as EL
-- @
--
-- Since: 0.4.5
--
-----------------------------------------------------------------------------

module Data.Enumerator.List (

	  head
	, drop
	, dropWhile
	, take
	, takeWhile
	, consume
	, require
	, isolate

	) where
import Data.Enumerator hiding (consume, head, peek, drop, dropWhile)
import Control.Exception (ErrorCall(..))
import Prelude hiding (head, drop, dropWhile, take, takeWhile)
import qualified Data.List as L


-- | Get the next element from the stream, or 'Nothing' if the stream has
-- ended.
--
-- Since: 0.4.5

head :: Monad m => Iteratee a m (Maybe a)
head = continue loop where
	loop (Chunks []) = head
	loop (Chunks (x:xs)) = yield (Just x) (Chunks xs)
	loop EOF = yield Nothing EOF


-- | @drop n@ ignores /n/ input elements from the stream.
--
-- Since: 0.4.5

drop :: Monad m => Integer -> Iteratee a m ()
drop n | n <= 0 = return ()
drop n = continue (loop n) where
	loop n' (Chunks xs) = iter where
		len = L.genericLength xs
		iter = if len < n'
			then drop (n' - len)
			else yield () (Chunks (L.genericDrop n' xs))
	loop _ EOF = yield () EOF


-- | @dropWhile p@ ignores input from the stream until the first element
-- which does not match the predicate.
--
-- Since: 0.4.5

dropWhile :: Monad m => (a -> Bool) -> Iteratee a m ()
dropWhile p = continue loop where
	loop (Chunks xs) = case L.dropWhile p xs of
		[] -> continue loop
		xs' -> yield () (Chunks xs')
	loop EOF = yield () EOF


-- | @take n@ extracts the next /n/ elements from the stream, as a list.
--
-- Since: 0.4.5

take :: Monad m => Integer -> Iteratee a m [a]
take n | n <= 0 = return []
take n = continue (loop id n) where
	len = L.genericLength
	loop acc n' (Chunks xs)
		| len xs < n' = continue (loop (acc . (xs ++)) (n' - len xs))
		| otherwise   = let
			(xs', extra) = L.genericSplitAt n' xs
			in yield (acc xs') (Chunks extra)
	loop acc _ EOF = yield (acc []) EOF


-- | @takeWhile p@ extracts input from the stream until the first element
-- which does not match the predicate.
--
-- Since: 0.4.5

takeWhile :: Monad m => (a -> Bool) -> Iteratee a m [a]
takeWhile p = continue (loop id) where
	loop acc (Chunks []) = continue (loop acc)
	loop acc (Chunks xs) = case Prelude.span p xs of
		(_, []) -> continue (loop (acc . (xs ++)))
		(xs', extra) -> yield (acc xs') (Chunks extra)
	loop acc EOF = yield (acc []) EOF


-- | Read all remaining input elements from the stream, and return as a list.
--
-- Since: 0.4.5

consume :: Monad m => Iteratee a m [a]
consume = continue (loop id) where
	loop acc (Chunks []) = continue (loop acc)
	loop acc (Chunks xs) = continue (loop (acc . (xs ++)))
	loop acc EOF = yield (acc []) EOF


-- | @require n@ buffers input until at least /n/ elements are available, or
-- throws an error if the stream ends early.
--
-- Since: 0.4.5

require :: Monad m => Integer -> Iteratee a m ()
require n | n <= 0 = return ()
require n = continue (loop id n) where
	len = L.genericLength
	loop acc n' (Chunks xs)
		| len xs < n' = continue (loop (acc . (xs ++)) (n' - len xs))
		| otherwise   = yield () (Chunks (acc xs))
	loop _ _ EOF = throwError (ErrorCall "require: Unexpected EOF")


-- | @isolate n@ reads at most /n/ elements from the stream, and passes them
-- to its iteratee. If the iteratee finishes early, elements continue to be
-- consumed from the outer stream until /n/ have been consumed.
--
-- Since: 0.4.5

isolate :: Monad m => Integer -> Enumeratee a a m b
isolate n step | n <= 0 = return step
isolate n (Continue k) = continue loop where
	len = L.genericLength
	
	loop (Chunks []) = continue loop
	loop (Chunks xs)
		| len xs <= n = k (Chunks xs) >>== isolate (n - len xs)
		| otherwise = let
			(s1, s2) = L.genericSplitAt n xs
			in k (Chunks s1) >>== (\step -> yield step (Chunks s2))
	loop EOF = k EOF >>== (\step -> yield step EOF)
isolate n step = drop n >> return step