(import (sx coroutines))

(defsuite
  "coroutine"
  (deftest
    "coroutine? recognizes coroutine objects"
    (let
      ((co (make-coroutine (fn () nil))))
      (assert (coroutine? co))
      (assert= false (coroutine? 42))
      (assert= false (coroutine? "hello"))
      (assert= false (coroutine? nil))
      (assert= false (coroutine? (list)))))
  (deftest
    "coroutine-alive? true for ready coroutine"
    (let
      ((co (make-coroutine (fn () nil))))
      (assert (coroutine-alive? co))))
  (deftest
    "coroutine-alive? false for non-coroutine"
    (assert= false (coroutine-alive? 42)))
  (deftest
    "immediate return — done true, value is body result"
    (let
      ((co (make-coroutine (fn () 42))))
      (let
        ((r (coroutine-resume co nil)))
        (assert= true (get r "done"))
        (assert= 42 (get r "value")))))
  (deftest
    "immediate nil return"
    (let
      ((co (make-coroutine (fn () nil))))
      (let
        ((r (coroutine-resume co nil)))
        (assert= true (get r "done"))
        (assert= nil (get r "value")))))
  (deftest
    "coroutine-alive? false after completion"
    (let
      ((co (make-coroutine (fn () nil))))
      (coroutine-resume co nil)
      (assert= false (coroutine-alive? co))))
  (deftest
    "single yield — done false on yield, done true on finish"
    (let
      ((co (make-coroutine (fn () (coroutine-yield 10) 20))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co nil)))
          (assert= false (get r1 "done"))
          (assert= 10 (get r1 "value"))
          (assert= true (get r2 "done"))
          (assert= 20 (get r2 "value"))))))
  (deftest
    "coroutine-alive? true between yield and next resume"
    (let
      ((co (make-coroutine (fn () (coroutine-yield nil) nil))))
      (assert (coroutine-alive? co))
      (coroutine-resume co nil)
      (assert (coroutine-alive? co))
      (coroutine-resume co nil)
      (assert= false (coroutine-alive? co))))
  (deftest
    "three yields then return"
    (let
      ((co (make-coroutine (fn () (coroutine-yield "a") (coroutine-yield "b") (coroutine-yield "c") "z"))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co nil)))
          (let
            ((r3 (coroutine-resume co nil)))
            (let
              ((r4 (coroutine-resume co nil)))
              (assert= "a" (get r1 "value"))
              (assert= false (get r1 "done"))
              (assert= "b" (get r2 "value"))
              (assert= false (get r2 "done"))
              (assert= "c" (get r3 "value"))
              (assert= false (get r3 "done"))
              (assert= "z" (get r4 "value"))
              (assert= true (get r4 "done"))))))))
  (deftest
    "final return vs yield — done flag distinguishes them"
    (let
      ((co (make-coroutine (fn () (coroutine-yield "yielded") "returned"))))
      (let
        ((y (coroutine-resume co nil)))
        (let
          ((r (coroutine-resume co nil)))
          (assert= false (get y "done"))
          (assert= "yielded" (get y "value"))
          (assert= true (get r "done"))
          (assert= "returned" (get r "value"))))))
  (deftest
    "resume val becomes yield return value"
    (let
      ((co (make-coroutine (fn () (let ((received (coroutine-yield "first"))) received)))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co 99)))
          (assert= "first" (get r1 "value"))
          (assert= false (get r1 "done"))
          (assert= 99 (get r2 "value"))
          (assert= true (get r2 "done"))))))
  (deftest
    "multiple resume values passed through yields"
    (let
      ((co (make-coroutine (fn () (let ((a (coroutine-yield 1))) (let ((b (coroutine-yield 2))) (+ a b)))))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co 10)))
          (let
            ((r3 (coroutine-resume co 20)))
            (assert= 1 (get r1 "value"))
            (assert= 2 (get r2 "value"))
            (assert= true (get r3 "done"))
            (assert= 30 (get r3 "value")))))))
  (deftest
    "coroutine captures lexical environment"
    (let
      ((x 10)
        (co
          (make-coroutine
            (fn () (coroutine-yield (* x 2)) (* x 3)))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co nil)))
          (assert= 20 (get r1 "value"))
          (assert= 30 (get r2 "value"))))))
  (deftest
    "resuming dead coroutine raises error"
    (let
      ((co (make-coroutine (fn () nil))))
      (coroutine-resume co nil)
      (assert-throws (fn () (coroutine-resume co nil)))))
  (deftest
    "coroutine drives iteration via recursive body"
    (let
      ((co (make-coroutine (fn () (define loop (fn (i) (when (< i 4) (coroutine-yield i) (loop (+ i 1))))) (loop 0))))
        (results (list)))
      (let
        drive
        ()
        (let
          ((r (coroutine-resume co nil)))
          (when
            (not (get r "done"))
            (append! results (get r "value"))
            (drive))))
      (assert= 4 (len results))
      (assert= 0 (nth results 0))
      (assert= 1 (nth results 1))
      (assert= 2 (nth results 2))
      (assert= 3 (nth results 3))))
  (deftest
    "nested coroutine — inner resumed from outer body"
    (let
      ((inner (make-coroutine (fn () (coroutine-yield "inner-a") "inner-done")))
        (outer
          (make-coroutine
            (fn
              ()
              (let
                ((i1 (coroutine-resume inner nil)))
                (coroutine-yield (get i1 "value")))
              (let ((i2 (coroutine-resume inner nil))) (get i2 "value"))))))
      (let
        ((o1 (coroutine-resume outer nil)))
        (let
          ((o2 (coroutine-resume outer nil)))
          (assert= false (get o1 "done"))
          (assert= "inner-a" (get o1 "value"))
          (assert= true (get o2 "done"))
          (assert= "inner-done" (get o2 "value"))))))
  (deftest
    "two independent coroutines interleave correctly"
    (let
      ((co1 (make-coroutine (fn () (coroutine-yield 1) 5)))
        (co2
          (make-coroutine (fn () (coroutine-yield 2) 6))))
      (let
        ((a (coroutine-resume co1 nil)))
        (let
          ((b (coroutine-resume co2 nil)))
          (let
            ((c (coroutine-resume co1 nil)))
            (let
              ((d (coroutine-resume co2 nil)))
              (assert= false (get a "done"))
              (assert= 1 (get a "value"))
              (assert= false (get b "done"))
              (assert= 2 (get b "value"))
              (assert= true (get c "done"))
              (assert= 5 (get c "value"))
              (assert= true (get d "done"))
              (assert= 6 (get d "value"))))))))
  (deftest
    "coroutine state field is ready before first resume"
    (let
      ((co (make-coroutine (fn () (coroutine-yield 1)))))
      (assert= "ready" (get co "state"))))
  (deftest
    "coroutine state field is suspended between yields"
    (let
      ((co (make-coroutine (fn () (coroutine-yield 1) 2))))
      (coroutine-resume co nil)
      (assert= "suspended" (get co "state"))))
  (deftest
    "coroutine state field is dead after completion"
    (let
      ((co (make-coroutine (fn () nil))))
      (coroutine-resume co nil)
      (assert= "dead" (get co "state"))))
  (deftest
    "yield works when called from nested helper function"
    (let
      ((co (make-coroutine (fn () (define helper (fn (x) (coroutine-yield x))) (helper 10) (helper 20)))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co nil)))
          (let
            ((r3 (coroutine-resume co nil)))
            (assert= false (get r1 "done"))
            (assert= 10 (get r1 "value"))
            (assert= false (get r2 "done"))
            (assert= 20 (get r2 "value"))
            (assert= true (get r3 "done")))))))
  (deftest
    "initial resume argument is ignored by ready coroutine"
    (let
      ((co (make-coroutine (fn () (coroutine-yield 42)))))
      (let
        ((r (coroutine-resume co "ignored")))
        (assert= false (get r "done"))
        (assert= 42 (get r "value")))))
  (deftest
    "coroutine with mutable closure state"
    (let
      ((counter {:value 0}))
      (let
        ((co (make-coroutine (fn () (dict-set! counter "value" 1) (coroutine-yield "a") (dict-set! counter "value" 2) (coroutine-yield "b")))))
        (assert= 0 (get counter "value"))
        (coroutine-resume co nil)
        (assert= 1 (get counter "value"))
        (coroutine-resume co nil)
        (assert= 2 (get counter "value")))))
  (deftest
    "coroutine can yield complex values"
    (let
      ((co (make-coroutine (fn () (coroutine-yield (list 1 2 3)) (coroutine-yield {:key "val"})))))
      (let
        ((r1 (coroutine-resume co nil)))
        (let
          ((r2 (coroutine-resume co nil)))
          (assert= false (get r1 "done"))
          (assert= 3 (len (get r1 "value")))
          (assert= false (get r2 "done"))
          (assert= "val" (get (get r2 "value") "key"))))))
  (deftest
    "round-robin scheduling of multiple coroutines"
    (let
      ((results (list))
        (co1
          (make-coroutine
            (fn () (coroutine-yield "a") (coroutine-yield "b"))))
        (co2
          (make-coroutine
            (fn () (coroutine-yield "c") (coroutine-yield "d")))))
      (append! results (get (coroutine-resume co1 nil) "value"))
      (append! results (get (coroutine-resume co2 nil) "value"))
      (append! results (get (coroutine-resume co1 nil) "value"))
      (append! results (get (coroutine-resume co2 nil) "value"))
      (assert= 4 (len results))
      (assert= "a" (nth results 0))
      (assert= "c" (nth results 1))
      (assert= "b" (nth results 2))
      (assert= "d" (nth results 3))))
  (deftest
    "coroutines created from same factory share no state"
    (let
      ((make-counter (fn (start) (make-coroutine (fn () (define loop (fn (n) (coroutine-yield n) (loop (+ n 1)))) (loop start))))))
      (let
        ((c1 (make-counter 0)) (c2 (make-counter 100)))
        (let
          ((a (get (coroutine-resume c1 nil) "value")))
          (let
            ((b (get (coroutine-resume c2 nil) "value")))
            (let
              ((c (get (coroutine-resume c1 nil) "value")))
              (let
                ((d (get (coroutine-resume c2 nil) "value")))
                (assert= 0 a)
                (assert= 100 b)
                (assert= 1 c)
                (assert= 101 d))))))))
  (deftest
    "resuming non-coroutine raises error"
    (assert-throws (fn () (coroutine-resume "not-a-coroutine" nil)))))
