;;; Arl Standard Library - Looping Macros

(module looping
  (export until do-list loop recur)

  (import control :refer (cond))
  (import functional :refer (map every?))
  (import list :refer (call))
  (import types :refer (symbol? list-or-pair? list? null?))
  (import sequences :refer (zip))

  ;;' @section Looping Macros

  ;;
  ;; do-list loops (defined in terms of while, which is a built-in special form)
  ;;

  ;;' @description Loop over seq binding var each iteration. Expands to while.
  ;;' @param binding Pair of (var sequence-expr)
  ;;' @param body Expressions evaluated for each element with var bound
  ;;' @examples
  ;;' (define total 0)
  ;;' (do-list (x '(1 2 3))
  ;;'   (set! total (+ total x)))
  ;;' total                       ; => 6
  ;;'
  ;;' (define result (list))
  ;;' (do-list (ch '("a" "b" "c"))
  ;;'   (set! result (append result (list ch))))
  ;;' result                      ; => ("a" "b" "c")
  ;;' @assert
  ;;' (define _test-total 0)
  ;;' (do-list (_x '(1 2 3)) (set! _test-total (+ _test-total _x)))
  ;;' (assert-equal 6 _test-total)
  ;;' @note Expands to a `while` loop over the sequence converted to a list.
  ;;' @seealso while, map, loop
  (defmacro do-list (binding . body)
    (begin
      (define var (car binding))
      (define seq (car (cdr binding)))
      (define remaining (gensym "for_remaining"))
      `(begin
         (define ,remaining (_as-list ,seq))
         (while (not (null? ,remaining))
           (define ,var (car ,remaining))
           ,@body
           (set! ,remaining (cdr ,remaining))))))

  ;;
  ;; loop/recur, clojure-style
  ;;

  ;;' @description Signal a recur back to the enclosing `loop`. Only valid inside a `loop` body.
  ;;' @param args New values for each loop variable (must match loop binding count)
  ;;' @examples
  ;;' (loop ((i 5) (acc 0))
  ;;'   (if (= i 0) acc (recur (- i 1) (+ acc i))))  ; => 15
  ;;' @assert
  ;;' (assert-equal 15 (loop ((i 5) (acc 0)) (if (= i 0) acc (recur (- i 1) (+ acc i)))))
  ;;' @seealso loop
  (defmacro recur (. args)
    (stop "recur can only be used inside loop"))

  ;; Helper: Check if a form should skip recur rewriting
  ;;' @description Check if a form should skip recur rewriting (quote, lambda, defmacro, loop).
  ;;' @internal
  (define skip-rewrite-form?
    (lambda (head)
      (if (symbol? head)
        (begin
          (define name (as.character head))
          (or (= name "quote")
              (= name "quasiquote")
              (= name "lambda")
              (= name "defmacro")
              (= name "loop")))
        #f)))

  ;; Helper: Validate binding pairs
  ;;' @description Validate that a binding pair is a list of length 2.
  ;;' @internal
  (define valid-binding-pair?
    (lambda (pair)
      (if (list-or-pair? pair)
        (isTRUE (= (length pair) 2))
        #f)))

  ;;' @description Validate that bindings is a list of valid binding pairs.
  ;;' @internal
  (define valid-bindings?
    (lambda (bindings)
      (if (list? bindings)
        (every? valid-binding-pair? bindings)
        #f)))

  ;; Helper: Transform recur calls in an expression
  ;;' @description Transform recur calls in an expression tree.
  ;;' @internal
  (define rewrite-recur
    (lambda (expr vars recur-flag recur-vals)
      (if (list? expr)
        (begin
          (define head (car expr))
          (cond
            ;; Handle recur call
            ((and (symbol? head) (= (as.character head) "recur"))
             (begin
               (define args (cdr expr))
               (if (not (= (length args) (length vars)))
                 (stop (sprintf "recur expects %d argument(s)" (length vars))))
               `(begin
                  (set! ,recur-flag #t)
                  (set! ,recur-vals (list ,@args))
                  (return #nil))))

            ;; Skip forms that create new scopes or shouldn't be transformed
            ((skip-rewrite-form? head) expr)

            ;; Recursively rewrite other list forms
            (else
             (call (cons head (map (lambda (e) (rewrite-recur e vars recur-flag recur-vals))
                                   (cdr expr)))))))
        expr)))

  ;; Helper: Generate variable rebinding code (set! so LHS is not evaluated)
  ;;' @description Generate code to rebind loop variables from recur-vals list.
  ;;' @internal
  (define generate-rebindings
    (lambda (vars recur-vals)
      (define index-op (as.symbol "[["))
      (define rebinds
        (map (lambda (pair)
               (define var (car pair))
               (define idx (car (cdr pair)))
               (define nth-expr (call (list index-op recur-vals idx)))
               `(set! ,var ,nth-expr))
             (zip vars (seq_len (length vars)))))
      (if (null? rebinds)
        (quote (begin))
        (if (= (length rebinds) 1)
          (car rebinds)
          `(begin ,@rebinds)))))

  ;;' @description Clojure-style loop/recur that expands to an R while loop.
  ;;' @param bindings List of (var init-value) pairs for loop variables
  ;;' @param body Loop body; use (recur ...) to iterate with new values
  ;;' @examples
  ;;' ;; Factorial via loop/recur
  ;;' (loop ((n 5) (acc 1))
  ;;'   (if (= n 0)
  ;;'     acc
  ;;'     (recur (- n 1) (* acc n))))  ; => 120
  ;;'
  ;;' ;; Sum of a list via loop/recur
  ;;' (loop ((xs '(1 2 3 4)) (total 0))
  ;;'   (if (null? xs)
  ;;'     total
  ;;'     (recur (cdr xs) (+ total (car xs)))))  ; => 10
  ;;' @assert
  ;;' (assert-equal 120 (loop ((n 5) (acc 1)) (if (= n 0) acc (recur (- n 1) (* acc n)))))
  ;;' (assert-equal 10 (loop ((xs '(1 2 3 4)) (total 0)) (if (null? xs) total (recur (cdr xs) (+ total (car xs))))))
  ;;' @note The compiler auto-optimizes self-tail-calls in `(define name (lambda ...))` patterns, so `loop`/`recur` is primarily useful for anonymous iteration and mutual recursion patterns. Expands to an R `while` loop internally.
  ;;' @seealso recur, while, do-list
  (defmacro loop (bindings . body)
    (begin
      ;; Validate bindings or fall back to runtime loop function
      (if (not (valid-bindings? bindings))
        `((lambda (loop-fn) (loop-fn ,bindings ,@body)) loop)
        (begin
          ;; Extract variables and generate gensyms
          (define vars (map car bindings))
          (define recur-flag (gensym "recur_flag"))
          (define recur-vals (gensym "recur_vals"))
          (define loop-result (gensym "loop_result"))
          (define loop-step (gensym "loop_step"))
          (define loop-done (gensym "loop_done"))

          ;; Rewrite body to transform recur calls
          (define rewritten-body
            (map (lambda (expr) (rewrite-recur expr vars recur-flag recur-vals))
                 body))

          ;; Generate the rebinding block
          (define rebind-block (generate-rebindings vars recur-vals))

          ;; Emit the loop structure (use done flag instead of break so we stay inside R's while)
          `(begin
             ;; Initialize loop variables
             ,@(map (lambda (pair) `(define ,(car pair) ,(car (cdr pair)))) bindings)
             (define ,recur-flag #f)
             (define ,recur-vals (list))
             (define ,loop-result #f)
             (define ,loop-done #f)
             (define ,loop-step (lambda () (begin ,@rewritten-body)))

             (while (not ,loop-done)
               (set! ,recur-flag #f)
               (set! ,loop-result (,loop-step))
               (if ,recur-flag
                 (r-eval ,rebind-block (current-env))
                 (set! ,loop-done #t)))

             ,loop-result)))))

  ;;
  ;; misc loops
  ;;

  ;;' @description Repeat body until test is truthy.
  ;;' @param test Condition checked before each iteration; loop stops when truthy
  ;;' @param body Expressions evaluated each iteration
  ;;' @examples
  ;;' (define i 0)
  ;;' (until (= i 5)
  ;;'   (set! i (+ i 1)))
  ;;' i                          ; => 5
  ;;' @assert
  ;;' (define _test-i 0)
  ;;' (until (= _test-i 5) (set! _test-i (+ _test-i 1)))
  ;;' (assert-equal 5 _test-i)
  ;;' @seealso while
  (defmacro until (test . body)
    `(while (not ,test) ,@body))
)
