-;; I want to use deftype, but literal deftype'd objects cannot be
-;; 9embedded in generated code.
+(deftype Context [before after])
 
-(defn TestCase [contexts locals children]
-  (with-meta {:contexts contexts
-              :locals locals
-              :children children}
-    {:type ::TestCase}))
+(deftype Assertion [locals form])
 
-(defn Context [before after]
-  (with-meta {:before before
-              :after after}
-    {:type ::Context}))
+(deftype TestCase [contexts children])
+
+(defprotocol TestCompilable
+  (test-compile [x]))
 
 (defprotocol TestSuccess
   (success? [r] "Returns true if r is a 100% successful result."))

  TestSuccess (success? [] false))

(defn format-assertion [a]
  `(try (if ~a
          (AssertionPassed '~a)
          (AssertionFailed '~a))
        (catch Throwable t#
          (AssertionThrown '~a t#))))

(defn format-tests [t]
  (let [active (gensym "active")]
    `(fn [~active ~@(:locals t) & states#]
       (lazy-seq
        ~(loop [r nil
                children (reverse (:children t))]
           (if (seq children)
             (let [c (first children)]
               (recur (if (= ::TestCase (type c))
                        `(cons (run-tests '~c ~active) ~r)
                        `(cons ~(format-assertion c) ~r))
                      (next children)))
             r))))))

(defn compile-tests [t]
  (eval (format-tests t)))

(defn unchunked-map [f coll]
  (lazy-seq

     (cons (f (first coll))
           (unchunked-map f (next coll))))))

(def mtest-compile (memoize test-compile))

(extend-protocol TestCompilable
 ::Assertion
 (test-compile [a] (eval (format-assertion a)))
 ::Context
 (test-compile [c] (fn ([] (when-let [f (:before c)] (f)))
                     ([state] (when-let [f (:after c)] (f state)))))
 ::TestCase
 (test-compile [t]
  (let [cfns (map mtest-compile (:contexts t))
        tfns (map mtest-compile (:children t))
        open (fn [active] (map (fn [f]
                                 (prn "Trying to open cfn" f)
                                 (or (active f) (f))) cfns))
        close (fn [active states]
                  (map (fn [c state]
                         (when-not (contains? active c) (c state)))
                       (reverse cfns) (reverse states)))]
    (fn [active & states]
      (prn "Active is" active)
       (let [states (open active)
             merged (merge active (zipmap cfns states))
             _ (prn "Merged is" merged)
             results (unchunked-map #(apply % merged states) tfns)]
         (close active states)
         (TestResult (str t) results))))))

(defn results-seq [r]
  (tree-seq :children :children r))

(def c1 (Context (fn [] (prn "Opening lazy context c1.") 1) nil))

(def a1 (Assertion '[a] '(do (prn "Executing assertion a1.")
                             (pos? a))))

(def t1 (TestCase [c1] [a1]))

(def t2 (TestCase [c1] [a1 a1 a1]))

(def t3 (TestCase [c1] [t1 t1 t1]))