[<AutoOpen>]
module FSharp.Monad.Iteratee.Core

// [snippet: Iteratee]
/// A stream of chunks of data generated by an Enumerator.
/// The stream can be composed of chunks of 'a, empty blocks indicating a wait, or an EOF marker.
/// In Haskell, the Chunk is usually composed of a list of ListLike type, but F# doesn't support
/// Monad Transforms or ^M in type declarations. Thus, the Chunk is left open to various internal
/// types, but a bit more work must be done in order to maintain the meaningfulness of "chunk".
/// That said, the 'a allows a large number of chunk-y types to be used, including other monads.
/// Be aware that when using #seq<_> types, you will need to check for both Seq.empty ([]) and Empty.
type Stream<'a> =
  | Chunk of 'a
  | Empty
  | EOF

/// The iteratee is a stream consumer that will consume a stream of data until either 
/// it receives an EOF or meets its own requirements for consuming data. The iteratee
/// will return Continue whenever it is ready to receive the next chunk. An iteratee
/// is fed data by an Enumerator, which generates a Stream. 
type Iteratee<'el,'acc> =
  | Continue of (Stream<'el> -> Iteratee<'el,'acc>)
  | Yield of 'acc * Stream<'el>
  | Error of exn

/// An enumerator generates a stream of data and feeds an iteratee, returning a new iteratee.
type Enumerator<'el,'acc> = Iteratee<'el,'acc> -> Iteratee<'el,'acc>

/// An Enumeratee is an Enumerator that feeds data streams to an internal iteratee.
type Enumeratee<'elo,'eli,'acc> = Iteratee<'eli,'acc> -> Iteratee<'elo, Iteratee<'eli,'acc>>

// TODO: Make calls to bind tail recursive.
let rec bind m f =
  match m with
  | Continue k -> Continue(fun s -> bind (k s) f)
  | Error e -> Error e
  | Yield(x, Empty) -> f x
  | Yield(x, extra) ->
      match f x with
      | Continue k -> k extra
      | Error e -> Error e
      | Yield(acc',_) -> Yield(acc', extra)

let combine comp1 comp2 =
  let binder () = comp2
  bind comp1 binder

type IterateeBuilder() =
  member this.Return(x) = Yield(x, Empty)
  member this.ReturnFrom(m:Iteratee<_,_>) = m
  member this.Bind(m, k) = bind m k
  member this.Zero() = Yield((), Empty)
  member this.Combine(comp1, comp2) = combine comp1 comp2
  member this.Delay(f) = bind (Yield((), Empty)) f
let iteratee = IterateeBuilder()
// [/snippet]
module Operators =
  let inline returnM x = Yield(x, Empty)
  let inline (>>=) m f = bind m f
  let inline (<*>) f m = bind f (fun f' -> bind m (fun m' -> returnM (f' m')))

// [snippet: Iteratee runners]
let rec enumEOF = function
  | Yield(x,_) -> Yield(x,EOF)
  | Error e -> Error e
  | Continue k ->
      match k EOF with
      | Continue _ -> failwith "enumEOF: divergent iteratee"
      | i -> enumEOF i

let run i =
  match enumEOF i with
  | Error e -> Choice1Of2 e
  | Yield(x,_) -> Choice2Of2 x
  | Continue _ -> failwith "run: divergent iteratee"

let run_ i =
  match run i with
  | Choice1Of2 e -> raise e
  | x -> x
// [/snippet]

module FSharp.Monad.Iteratee.List

open System
open Operators

(* ========= Extensions ========= *)

module List =

  let split pred l =
    let rec loop l cont =
      match l with
      | [] -> ([],[])
      | x::[] when not (pred x) -> (cont l, [])
      | x::xs when pred x -> (cont [], l)
      | x::xs when not (pred x) -> loop xs (fun rest -> cont (x::rest))
      | _ -> failwith "List.split: Unrecognized pattern"
    loop l id
  
  let splitAt n l =
    let pred i = i >= n
    let rec loop i l cont =
      match l with
      | [] -> ([],[])
      | x::[] when not (pred i) -> (cont l, [])
      | x::xs when pred i -> (cont [], l)
      | x::xs when not (pred i) -> loop (i+1) xs (fun rest -> cont (x::rest))
      | _ -> failwith "List.splitAt: Unrecognized pattern"
    loop 0 l id

(* ========= Iteratees ========= *)
// [snippet: Sample iteratees]
let length<'a> : Iteratee<'a list, int> =
  let rec step n = function
    | Empty | Chunk [] -> Continue (step n)
    | Chunk x          -> Continue (step (n + 1))
    | EOF         as s -> Yield(n, s)
  Continue (step 0)

let rec peek =
  let rec step = function
    | Empty | Chunk []  -> peek
    | Chunk(x::xs) as s -> Yield(Some x, s)
    | s                 -> Yield(None, s)
  Continue step

let rec head =
  let rec step = function
    | Empty | Chunk [] -> head
    | Chunk(x::xs)     -> Yield(Some x, (Chunk xs))
    | EOF              -> Yield(None, EOF)
  Continue step

let rec drop n =
  let rec step = function
    | Empty | Chunk [] -> Continue step
    | Chunk x          -> drop (n - 1)
    | EOF         as s -> Yield((), s)
  if n = 0 then Yield((), Empty) else Continue step

let split (pred:char -> bool) =
  let rec step before = function
    | Empty | Chunk [] -> Continue (step before)
    | Chunk str ->
        match List.split pred str with
        | (_,[]) -> Continue (step (before @ str))
        | (str,tail) -> Yield((before @ str), Chunk tail)
    | s -> Yield(before, s)
  Continue (step [])

let heads str =
  let rec loop count str =
    match count, str with
    | (count, []) -> Yield(count, EOF)
    | (count, str) -> Continue (step count str)
  and step count str s =
    let str = List.ofSeq str
    match str, s with
    | str, Empty -> loop count str
    | str, (Chunk []) -> loop count str
    | c::t, (Chunk (c'::t')) ->
        if c = c' then step (count + 1) t (Chunk t') 
        else Yield(count, (Chunk (c'::t')))
    | _, s -> Yield(count, s)
  loop 0 str

let readLines =
  let toString chars = String(Array.ofList chars)
  let crlf = ['\r';'\n']
  let lf = ['\n']
  let isNewline c = c = '\r' || c = '\n'
  let terminators = heads crlf >>= fun n -> if n = 0 then heads lf else Yield(n, Empty)
  let rec lines acc = split isNewline >>= fun l -> terminators >>= check acc l
  and check acc l count =
    match l, count with
    | [], _ -> Yield (Choice2Of2 (List.rev acc |> List.map toString), EOF)
    | l, 0 -> Yield (Choice1Of2 (List.rev acc |> List.map toString), Chunk l)
    | l, _ -> lines (l::acc)
  lines []
// [/snippet]

(* ========= Enumerators ========= *)
// [snippet: Sample enumerators]
//val enumerate :: 'a list -> Enumerator<'a list,'b>
let rec enumerate input = fun i ->
  match input, i with
  | [], Continue k -> Continue k
  | (x::xs), Continue k -> enumerate xs (k (Chunk [x]))
  | _, i -> i

// val enumeratePure1Chunk :: 'a list -> Enumerator<'a list,'b>
let enumeratePure1Chunk (str:'a list) = function
  | Continue k -> k (Chunk str)
  | i -> i

// val enumeratePureNChunk :: 'a list -> int -> Enumerator<'a list,'b>
let rec enumeratePureNChunk str n = fun i ->
  match str, i with
  | _::_, Continue k ->
      let (s1, s2) = List.splitAt n str
      enumeratePureNChunk s2 n (k (Chunk s1))
  | _, i -> i
// [/snippet]