gleam_stdlib

import gleam/list

// Internal private representation of an Iterator
type Action(element) {
  Stop
  Continue(element, fn() -> Action(element))
}

/// An iterator is a lazily evaluated sequence of element.
///
/// Iterators are useful when working with collections that are too large to
/// fit in memory (or those that are infinite in size) as they only require the
/// elements currently being processed to be in memory.
///
/// As a lazy data structure no work is done when an iterator is filters,
/// mapped, etc, instead a new iterator is returned with these transformations
/// applied to the stream. Once the stream has all the required transformations
/// applied it can be evaluated using functions such as `fold` and `take`.
///
pub opaque type Iterator(element) {
  Iterator(continuation: fn() -> Action(element))
}

// Public API for iteration
pub type Step(element, accumulator) {
  Next(element: element, accumulator: accumulator)
  Done
}

// Creating Iterators
fn do_unfold(initial, f) {
  fn() {
    case f(initial) {
      Next(x, acc) -> Continue(x, do_unfold(acc, f))
      Done -> Stop
    }
  }
}

/// Create an iterator from a given function and accumulator.
///
/// The function is called on the accumulator and return either `Done`,
/// indicating the iterator has no more elements, or `Next` which contains a
/// new element and accumulator. The element is yielded by the iterator and the
/// new accumulator is used with the function to compute the next element in
/// the sequence.
///
/// ## Examples
///
///    > unfold(from: 5, with: fn(n) {
///    >  case n {
///    >    0 -> Done
///    >    n -> Next(element: n, accumulator: n - 1)
///    >  }
///    > })
///    > |> to_list
///    [5, 4, 3, 2, 1]
///
pub fn unfold(
  from initial: acc,
  with f: fn(acc) -> Step(element, acc),
) -> Iterator(element) {
  initial
  |> do_unfold(f)
  |> Iterator
}

// TODO: test
/// Create an iterator that yields values created by calling a given function
/// repeatedly.
///
pub fn repeatedly(f: fn() -> element) -> Iterator(element) {
  unfold(Nil, fn(acc) { Next(f(), acc) })
}

/// Create an iterator that returns the same value infinitely.
///
/// ## Examples
///
///    > repeat(10)
///    > |> take(4)
///    > |> to_list
///    [10, 10, 10, 10]
///
pub fn repeat(x: element) -> Iterator(element) {
  repeatedly(fn() { x })
}

/// Create an iterator the yields each element in a given list.
///
/// ## Examples
///
///    > from_list([1, 2, 3, 4]) |> to_list
///    [1, 2, 3, 4]
///
pub fn from_list(list: List(element)) -> Iterator(element) {
  let yield = fn(acc) {
    case acc {
      [] -> Done
      [head, ..tail] -> Next(head, tail)
    }
  }
  unfold(list, yield)
}

// Consuming Iterators
fn do_fold(iterator, initial, f) {
  case iterator() {
    Continue(element, iterator) -> do_fold(iterator, f(element, initial), f)
    Stop -> initial
  }
}

/// Reduce an iterator of elements into a single value by calling a given
/// function on each element in turn.
///
/// If called on an iterator of infinite length then this function will never
/// return.
///
/// If you do not care about the end value and only wish to evaluate the
/// iterator for side effects consider using the `run` function instead.
///
/// ## Examples
///
///    > [1, 2, 3, 4]
///    > |> from_list
///    > |> fold(from: 0, with: fn(element, acc) { element + acc })
///    10
///
pub fn fold(
  over iterator: Iterator(e),
  from initial: acc,
  with f: fn(e, acc) -> acc,
) -> acc {
  iterator.continuation
  |> do_fold(initial, f)
}

// TODO: test
/// Evaluate all elements in a given stream. This function is useful for when
/// you wish to trigger any side effects that would occur when evaluating
/// the iterator.
///
pub fn run(iterator) -> Nil {
  fold(iterator, Nil, fn(_, acc) { acc })
}

/// Evaluate an iterator and return all the elements as a list.
///
/// If called on an iterator of infinite length then this function will never
/// return.
///
/// ## Examples
///
///   > [1, 2, 3] |> from_list |> map(fn(x) { x * 2 }) |> to_list
///   [2, 4, 6]
///
pub fn to_list(iterator: Iterator(element)) -> List(element) {
  iterator
  |> fold([], fn(e, acc) { [e, ..acc] })
  |> list.reverse
}

fn do_take(iterator, desired, acc) {
  case desired > 0 {
    True -> case iterator() {
      Continue(
        element,
        iterator,
      ) -> do_take(iterator, desired - 1, [element, ..acc])
      Stop -> acc
    }
    False -> acc
  }
}

/// Evaluate a desired number of elements from an iterator and return them in a
/// list.
///
/// If the iterator does not have enough elements all of them are returned.
///
/// ## Examples
///
///    > [1, 2, 3, 4, 5] |> from_list |> take(up_to: 3)
///    [1, 2, 3]
///
///    > [1, 2] |> from_list |> take(up_to: 3)
///    [1, 2]
///
pub fn take(from iterator: Iterator(e), up_to desired: Int) -> List(e) {
  iterator.continuation
  |> do_take(desired, [])
  |> list.reverse
}

fn do_drop(iterator, desired) {
  case desired > 0 {
    True -> case iterator() {
      Continue(_, iterator) -> do_drop(iterator, desired - 1)
      Stop -> fn() { Stop }
    }
    False -> iterator
  }
}

/// Evaluate and discard the first N elements in an iterator, returning a new
/// iterator.
///
/// If the iterator does not have enough elements an empty iterator is
/// returned.
///
/// This function does not evaluate the elements of the iterator, the
/// computation is performed when the iterator is later run.
///
/// ## Examples
///
///    > [1, 2, 3, 4, 5] |> from_list |> drop(up_to: 3) |> to_list
///    [4, 5]
///
///    > [1, 2] |> from_list |> drop(up_to: 3) |> to_list
///    []
///
pub fn drop(from iterator: Iterator(e), up_to desired: Int) -> Iterator(e) {
  iterator.continuation
  |> do_drop(desired)
  |> Iterator
}

fn do_map(continuation, f) {
  fn() {
    case continuation() {
      Continue(e, continuation) -> Continue(f(e), do_map(continuation, f))
      Stop -> Stop
    }
  }
}

/// Create an iterator from an existing iterator and a transformation function.
///
/// Each element in the new iterator will be the result of calling the given
/// function on the elements in the given iterator.
///
/// This function does not evaluate the elements of the iterator, the
/// computation is performed when the iterator is later run.
///
/// ## Examples
///
///    > [1, 2, 3] |> from_list |> map(fn(x) { x * 2 }) |> to_list
///    [2, 4, 6]
///
pub fn map(over iterator: Iterator(a), with f: fn(a) -> b) -> Iterator(b) {
  iterator.continuation
  |> do_map(f)
  |> Iterator
}

fn do_filter(iterator, predicate) {
  fn() {
    case iterator() {
      Continue(e, iterator) -> case predicate(e) {
        True -> Continue(e, do_filter(iterator, predicate))
        False -> do_filter(iterator, predicate)()
      }
      Stop -> Stop
    }
  }
}

/// Create an iterator from an existing iterator and a predicate function.
///
/// The new iterator will contain elements from the first iterator for which
/// the given function returns `True`.
///
/// This function does not evaluate the elements of the iterator, the
/// computation is performed when the iterator is later run.
///
/// ## Examples
///
///    > import gleam/int
///    > [1, 2, 3, 4] |> from_list |> filter(int.is_even) |> to_list
///    [2, 4]
///
pub fn filter(
  iterator: Iterator(a),
  for predicate: fn(a) -> Bool,
) -> Iterator(a) {
  iterator.continuation
  |> do_filter(predicate)
  |> Iterator
}

fn do_cycle(next: fn() -> Action(a), reset: fn() -> Action(a)) {
  fn() {
    case next() {
      Continue(e, iterator) -> Continue(e, do_cycle(iterator, reset))
      Stop -> do_cycle(reset, reset)()
    }
  }
}

/// Create an iterator that repeats a given iterator infinitely.
///
/// ## Examples
///
///    > [1, 2] |> from_list |> cycle |> take(6)
///    [1, 2, 1, 2, 1, 2]
///
pub fn cycle(iterator: Iterator(a)) -> Iterator(a) {
  iterator.continuation
  |> do_cycle(iterator.continuation)
  |> Iterator
}

fn do_range(current, limit, inc) -> fn() -> Action(Int) {
  case current == limit {
    True -> fn() { Stop }
    False -> fn() { Continue(current, do_range(current + inc, limit, inc)) }
  }
}

/// Create an iterator of ints, starting at a given start int and stepping by
/// one to a given end int.
///
/// ## Examples
///
///    > range(from: 1, to: 5) |> to_list
///    [1, 2, 3, 4]
///
///    > range(from: 1, to: -2) |> to_list
///    [1, 0, -1]
///
///    > range(from: 0, to: 0) |> to_list
///    []
///
pub fn range(from start: Int, to stop: Int) -> Iterator(Int) {
  case start < stop {
    True -> 1
    False -> -1
  }
  |> do_range(start, stop, _)
  |> Iterator
}