(;module:
  lux
  (lux (control functor
                monad
                comonad
                [cont #+ pending Cont]
                ["p" parser])
       [macro #+ with-gensyms]
       (macro ["s" syntax #+ syntax: Syntax])
       (data (coll [list "List/" Monad<List>])
             bool)))

## [Types]
(type: #export (Stream a)
  {#;doc "An infinite stream of values."}
  (Cont [a (Stream a)]))

## [Utils]
(def: (cycle' x xs init full)
  (All [a]
    (-> a (List a) a (List a) (Stream a)))
  (case xs
    #;Nil           (pending [x (cycle' init full init full)])
    (#;Cons x' xs') (pending [x (cycle' x' xs' init full)])))

## [Functions]
(def: #export (iterate f x)
  {#;doc "Create a stream by applying a function to a value, and to its result, on and on..."}
  (All [a]
    (-> (-> a a) a (Stream a)))
  (pending [x (iterate f (f x))]))

(def: #export (repeat x)
  {#;doc "Repeat a value forever."}
  (All [a]
    (-> a (Stream a)))
  (pending [x (repeat x)]))

(def: #export (cycle xs)
  {#;doc "Go over the elements of a list forever.

          The list should not be empty."}
  (All [a]
    (-> (List a) (Maybe (Stream a))))
  (case xs
    #;Nil          #;None
    (#;Cons x xs') (#;Some (cycle' x xs' x xs'))))

(do-template [<name> <return> <part>]
  [(def: #export (<name> s)
     (All [a] (-> (Stream a) <return>))
     (let [[h t] (cont;run s)]
       <part>))]

  [head a          h]
  [tail (Stream a) t])

(def: #export (nth idx s)
  (All [a] (-> Nat (Stream a) a))
  (let [[h t] (cont;run s)]
    (if (n.> +0 idx)
      (nth (n.dec idx) t)
      h)))

(do-template [<taker> <dropper> <splitter> <pred-type> <pred-test> <pred-step>]
  [(def: #export (<taker> pred xs)
     (All [a]
       (-> <pred-type> (Stream a) (List a)))
     (let [[x xs'] (cont;run xs)]
       (if <pred-test>
         (list& x (<taker> <pred-step> xs'))
         (list))))

   (def: #export (<dropper> pred xs)
     (All [a]
       (-> <pred-type> (Stream a) (Stream a)))
     (let [[x xs'] (cont;run xs)]
       (if <pred-test>
         (<dropper> <pred-step> xs')
         xs)))

   (def: #export (<splitter> pred xs)
     (All [a]
       (-> <pred-type> (Stream a) [(List a) (Stream a)]))
     (let [[x xs'] (cont;run xs)]
       (if <pred-test>
         (let [[tail next] (<splitter> <pred-step> xs')]
           [(#;Cons [x tail]) next])
         [(list) xs])))]
  
  [take-while drop-while split-while (-> a Bool) (pred x)      pred]
  [take       drop       split       Nat         (n.> +0 pred) (n.dec pred)]
  )

(def: #export (unfold step init)
  {#;doc "A stateful way of infinitely calculating the values of a stream."}
  (All [a b]
    (-> (-> a [a b]) a (Stream b)))
  (let [[next x] (step init)]
    (pending [x (unfold step next)])))

(def: #export (filter p xs)
  (All [a] (-> (-> a Bool) (Stream a) (Stream a)))
  (let [[x xs'] (cont;run xs)]
    (if (p x)
      (pending [x (filter p xs')])
      (filter p xs'))))

(def: #export (partition p xs)
  {#;doc "Split a stream in two based on a predicate.

          The left side contains all entries for which the predicate is true.

          The right side contains all entries for which the predicate is false."}
  (All [a] (-> (-> a Bool) (Stream a) [(Stream a) (Stream a)]))
  [(filter p xs) (filter (complement p) xs)])

## [Structures]
(struct: #export _ (Functor Stream)
  (def: (map f fa)
    (let [[h t] (cont;run fa)]
      (pending [(f h) (map f t)]))))

(struct: #export _ (CoMonad Stream)
  (def: functor Functor<Stream>)
  (def: unwrap head)
  (def: (split wa)
    (let [[head tail] (cont;run wa)]
      (pending [wa (split tail)]))))

## [Pattern-matching]
(syntax: #export (^stream& [patterns (s;form (p;many s;any))] body [branches (p;some s;any)])
  {#;doc (doc "Allows destructuring of streams in pattern-matching expressions."
              "Caveat emptor: Only use it for destructuring, and not for testing values within the streams."
              (let [(^stream& x y z _tail) (some-stream-func 1 2 3)]
                (func x y z)))}
  (with-gensyms [g!s]
    (let [body+ (` (let [(~@ (List/join (List/map (function [pattern]
                                                    (list (` [(~ pattern) (~ g!s)])
                                                          (` (cont;run (~ g!s)))))
                                                  patterns)))]
                     (~ body)))]
      (wrap (list& g!s body+ branches)))))