{-# LANGUAGE Rank2Types #-}

-- Converting from Enumerators to streams. The end goal is to `zip' two
-- enumerators

module IterateeZ where

import System.Posix
import Foreign.C
import Foreign.Ptr
import Foreign.Marshal.Alloc
import Prelude hiding (head, drop, dropWhile, take, break, catch)

import Iteratee
import Control.Monad.Identity
import LowLevelIO

-- ------------------------------------------------------------------------
-- First, we introduce `open enumerators' as a conservative extension
-- of the old enumerators. Open enumerators take the additional argument
-- cifd. If we pass check_if_doneM as the value of that argument, we
-- obtain old enumerators.

-- Enumerator in the `open' style
type EnumeratorO m a = 
 (EnumeratorM m a -> IterV a -> m (Iteratee a)) -> EnumeratorM m a



-- The pure n-chunk enumerator in the `open' style
-- It passes a given string to the iteratee in chunks no larger than n.
-- This enumerator does no IO and is useful for testing of base parsing
-- and handling of chunk boundaries
enum_pure_nchunk_open :: Monad m => String -> Int -> EnumeratorO m a
enum_pure_nchunk_open ""  _ cifd iter = return iter
enum_pure_nchunk_open str n cifd iter = cifd (enum_pure_nchunk_open s2 n cifd)$ 
					runIter iter (Chunk s1)
 where (s1,s2) = splitAt n str


-- when check_if_doneM is passed as the cifd argument, the recover the
-- original enum_pure_nchunk enumerator.
testpo2 =
    let (Right lines, rest)
	    = run $ runIdentity $ 
	      enum_pure_nchunk_open test_str1 5 check_if_doneM read_lines_rest
    in
    lines == ["header1: v1","header2: v2","header3: v3","header4: v4",
	      "header5: v5","header6: v6","header7: v7"]
    && rest == "rest\n"

-- The open enumerator to read from a file descriptor
enum_fd_open :: Fd -> EnumeratorO IO a
enum_fd_open fd cifd iter = allocaBytes (fromIntegral buffer_size) (loop iter)
 where
  buffer_size = 5 -- for tests; in real life, there should be 1024 or so
  loop iter p = do
   n <- myfdRead fd p buffer_size
   case n of
    Left errno -> return $ enum_err "IO error" iter
    Right 0 -> return iter
    Right n -> do
	 str <- peekCAStringLen (p,fromIntegral n)
	 cifd (\i -> loop i p) $ runIter iter (Chunk str)

-- Again, if we pass check_if_doneM as the cifd argument, the obtain
-- the original enum_fd. The tests below confirm that.

test_driver1 filepath = do
  fd <- openFd filepath ReadOnly Nothing defaultFileFlags
  putStrLn "About to read headers"
  result <- fmap run $ enum_fd_open fd check_if_doneM 
	                            read_lines_and_one_more_line
  closeFd fd
  putStrLn "Finished reading headers"
  case result of
   (Right headers,after,_) ->
       do
       putStrLn $ "The line after headers is: " ++ show after
       putStrLn "Complete headers"
       print headers
   (Left headers,_,status) ->
       do
       putStrLn $ "Problem " ++ show status
       putStrLn "Incomplete headers"
       print headers
 where
  read_lines_and_one_more_line = do
     lines  <- read_lines
     after  <- break (\c -> c == '\r' || c == '\n')
     status <- is_finished
     return (lines,after,status)


testz11 = test_driver1 "test1.txt" -- Complete headers, up to "header7: v7"
testz12 = test_driver1 "test2.txt" -- The same
testz13 = test_driver1 "test3.txt" -- "header3: v3", then EOF
testz14 = test_driver1 "/dev/null" -- Incomplete headers [], EOF


-- ------------------------------------------------------------------------
-- Converting from an enumerator to a stream

-- Like Stream, but recursive
data RStream m = RS_eof (Maybe ErrMsg) | RS_cons String (m (RStream m))

enumo_to_stream :: Monad m => 
		   (forall a. EnumeratorO m a) -> m (RStream m)
enumo_to_stream enumo = enumo cifd (iter r0) >>= return . project
 where
 -- a particular iteratee
 iter :: Monad m => RStream m -> Iteratee (RStream m)
 iter rstr = Iteratee (step rstr)
 step rstr (Chunk "")  = IE_cont (iter rstr) Nothing
 step rstr (Chunk str) = IE_done (RS_cons str (return rstr)) (Chunk "")
 step rstr (EOF e)     = IE_done rstr (EOF e)

 project :: Iteratee (RStream m) -> RStream m
 project iter = case runIter iter (EOF Nothing) of 
	        IE_done x _ -> x
	        IE_cont _ e -> RS_eof e

 cifd k (IE_done rstr _)        = return $ iter (ccat rstr k)
 cifd k (IE_cont x Nothing)  = k x
 cifd _ (IE_cont _ (Just e)) = return $ throwErr e


 ccat :: Monad m => RStream m -> EnumeratorM m (RStream m) -> RStream m
 ccat (RS_cons h t) k = RS_cons h (t >>= \tr -> ccatM tr k)
 ccatM (RS_eof Nothing) k = k (iter r0) >>= return . project
 ccatM r@RS_cons{} k = return $ ccat r k

 r0 = RS_eof Nothing


read_all_stream :: Monad m => RStream m -> m [String]
read_all_stream (RS_eof Nothing)  = return []
read_all_stream (RS_eof (Just e)) = error $ "read_all_stream: " ++ e
read_all_stream (RS_cons h t) = t >>= read_all_stream >>= return . (h:)

testpo3 =
 let str1 =runIdentity $ enumo_to_stream (enum_pure_nchunk_open test_str1 5) >>=
                         read_all_stream
 in concat str1 == test_str1

test_driver2 filepath = do
  fd <- openFd filepath ReadOnly Nothing defaultFileFlags
  putStrLn "About to read all"
  result <- enumo_to_stream (enum_fd_open fd) >>= read_all_stream
  closeFd fd
  print result

testz21 = test_driver2 "test1.txt" -- Complete headers, up to "header7: v7"
testz22 = test_driver2 "test2.txt" -- The same
testz23 = test_driver2 "test3.txt" -- "header3: v3", then EOF
testz24 = test_driver2 "/dev/null" -- Incomplete headers [], EOF