;; ==========================================================================
;; test-vm.sx — Tests for the bytecode VM (spec/vm.sx)
;;
;; Requires: test-framework.sx, compiler.sx, vm.sx loaded.
;; Tests the compile → bytecode → VM execution pipeline.
;; ==========================================================================


;; Helper: compile an SX expression and execute it on the VM.
;; Returns the result value.
(define vm-eval
  (fn (expr)
    (let ((code (compile expr)))
      (vm-execute-module
        (code-from-value code)
        {}))))

;; Helper: compile and run with a pre-populated globals dict.
(define vm-eval-with
  (fn (expr globals)
    (let ((code (compile expr)))
      (vm-execute-module (code-from-value code) globals))))


;; --------------------------------------------------------------------------
;; Constants and literals
;; --------------------------------------------------------------------------

(defsuite "vm-constants"
  (deftest "number constant"
    (assert-equal 42 (vm-eval 42)))

  (deftest "string constant"
    (assert-equal "hello" (vm-eval "hello")))

  (deftest "boolean true"
    (assert-equal true (vm-eval true)))

  (deftest "boolean false"
    (assert-equal false (vm-eval false)))

  (deftest "nil constant"
    (assert-nil (vm-eval nil)))

  (deftest "negative number"
    (assert-equal -7 (vm-eval -7)))

  (deftest "float constant"
    (assert-equal 3.14 (vm-eval 3.14))))


;; --------------------------------------------------------------------------
;; Arithmetic via primitives
;; --------------------------------------------------------------------------

(defsuite "vm-arithmetic"
  (deftest "addition"
    (assert-equal 5 (vm-eval '(+ 2 3))))

  (deftest "subtraction"
    (assert-equal 7 (vm-eval '(- 10 3))))

  (deftest "multiplication"
    (assert-equal 24 (vm-eval '(* 6 4))))

  (deftest "division"
    (assert-equal 5 (vm-eval '(/ 10 2))))

  (deftest "nested arithmetic"
    (assert-equal 14 (vm-eval '(+ (* 3 4) 2))))

  (deftest "three-arg addition"
    (assert-equal 15 (vm-eval '(+ 5 4 6)))))


;; --------------------------------------------------------------------------
;; Comparison and logic
;; --------------------------------------------------------------------------

(defsuite "vm-comparison"
  (deftest "equal numbers"
    (assert-equal true (vm-eval '(= 1 1))))

  (deftest "unequal numbers"
    (assert-equal false (vm-eval '(= 1 2))))

  (deftest "less than"
    (assert-equal true (vm-eval '(< 1 2))))

  (deftest "greater than"
    (assert-equal true (vm-eval '(> 5 3))))

  (deftest "not"
    (assert-equal true (vm-eval '(not false))))

  (deftest "not truthy"
    (assert-equal false (vm-eval '(not 42)))))


;; --------------------------------------------------------------------------
;; Control flow — if, when, cond, and, or
;; --------------------------------------------------------------------------

(defsuite "vm-control-flow"
  (deftest "if true branch"
    (assert-equal 1 (vm-eval '(if true 1 2))))

  (deftest "if false branch"
    (assert-equal 2 (vm-eval '(if false 1 2))))

  (deftest "if without else returns nil"
    (assert-nil (vm-eval '(if false 1))))

  (deftest "when true evaluates body"
    (assert-equal 42 (vm-eval '(when true 42))))

  (deftest "when false returns nil"
    (assert-nil (vm-eval '(when false 42))))

  (deftest "and short-circuits on false"
    (assert-equal false (vm-eval '(and true false 42))))

  (deftest "and returns last truthy"
    (assert-equal 3 (vm-eval '(and 1 2 3))))

  (deftest "or short-circuits on true"
    (assert-equal 1 (vm-eval '(or 1 false 2))))

  (deftest "or returns false when all falsy"
    (assert-equal false (vm-eval '(or false false false))))

  (deftest "cond first match"
    (assert-equal "one" (vm-eval '(cond (= 1 1) "one" (= 2 2) "two"))))

  (deftest "cond else clause"
    (assert-equal "none" (vm-eval '(cond (= 1 2) "one" :else "none"))))

  (deftest "case match"
    (assert-equal "two" (vm-eval '(case 2 1 "one" 2 "two" :else "other"))))

  (deftest "case else"
    (assert-equal "other" (vm-eval '(case 99 1 "one" 2 "two" :else "other")))))


;; --------------------------------------------------------------------------
;; Let bindings
;; --------------------------------------------------------------------------

(defsuite "vm-let"
  (deftest "single binding"
    (assert-equal 10 (vm-eval '(let ((x 10)) x))))

  (deftest "multiple bindings"
    (assert-equal 30 (vm-eval '(let ((x 10) (y 20)) (+ x y)))))

  (deftest "bindings are sequential"
    (assert-equal 11 (vm-eval '(let ((x 10) (y (+ x 1))) y))))

  (deftest "nested let"
    (assert-equal 3 (vm-eval '(let ((x 1)) (let ((y 2)) (+ x y))))))

  (deftest "inner let shadows outer"
    (assert-equal 99 (vm-eval '(let ((x 1)) (let ((x 99)) x)))))

  (deftest "let body returns last expression"
    (assert-equal 3 (vm-eval '(let ((x 1)) 1 2 3)))))


;; --------------------------------------------------------------------------
;; Lambda and function calls
;; --------------------------------------------------------------------------

(defsuite "vm-lambda"
  (deftest "lambda call"
    (assert-equal 7 (vm-eval '(let ((f (fn (x) (+ x 2)))) (f 5)))))

  (deftest "lambda with multiple params"
    (assert-equal 11 (vm-eval '(let ((add (fn (a b) (+ a b)))) (add 5 6)))))

  (deftest "higher-order: pass lambda to lambda"
    (assert-equal 10
      (vm-eval '(let ((apply-fn (fn (f x) (f x)))
                      (double (fn (n) (* n 2))))
                  (apply-fn double 5)))))

  (deftest "lambda returns lambda"
    (assert-equal 15
      (vm-eval '(let ((make-adder (fn (n) (fn (x) (+ n x)))))
                  (let ((add10 (make-adder 10)))
                    (add10 5))))))

  (deftest "immediately invoked lambda"
    (assert-equal 42 (vm-eval '((fn (x) (* x 2)) 21)))))


;; --------------------------------------------------------------------------
;; Closures and upvalues
;; --------------------------------------------------------------------------

(defsuite "vm-closures"
  (deftest "closure captures local"
    (assert-equal 10
      (vm-eval '(let ((x 10))
                  (let ((f (fn () x)))
                    (f))))))

  (deftest "closure captures through two levels"
    (assert-equal 30
      (vm-eval '(let ((x 10))
                  (let ((y 20))
                    (let ((f (fn () (+ x y))))
                      (f)))))))

  (deftest "two closures share upvalue"
    (assert-equal 42
      (vm-eval '(let ((x 0))
                  (let ((set-x (fn (v) (set! x v)))
                        (get-x (fn () x)))
                    (set-x 42)
                    (get-x))))))

  (deftest "closure mutation visible to sibling"
    (assert-equal 3
      (vm-eval '(let ((counter 0))
                  (let ((inc! (fn () (set! counter (+ counter 1)))))
                    (inc!)
                    (inc!)
                    (inc!)
                    counter))))))


;; --------------------------------------------------------------------------
;; Tail call optimization
;; --------------------------------------------------------------------------

(defsuite "vm-tco"
  (deftest "tail-recursive loop doesn't overflow"
    ;; Count down from 10000 — would overflow without TCO
    (assert-equal 0
      (vm-eval '(let ((loop (fn (n)
                              (if (<= n 0) 0
                                (loop (- n 1))))))
                  (loop 10000)))))

  (deftest "tail-recursive accumulator"
    (assert-equal 5050
      (vm-eval '(let ((sum (fn (n acc)
                             (if (<= n 0) acc
                               (sum (- n 1) (+ acc n))))))
                  (sum 100 0))))))


;; --------------------------------------------------------------------------
;; Collections
;; --------------------------------------------------------------------------

(defsuite "vm-collections"
  (deftest "list construction"
    (assert-equal (list 1 2 3) (vm-eval '(list 1 2 3))))

  (deftest "empty list"
    (assert-equal (list) (vm-eval '(list))))

  (deftest "dict construction"
    (let ((d (vm-eval '{:a 1 :b 2})))
      (assert-equal 1 (get d "a"))
      (assert-equal 2 (get d "b"))))

  (deftest "list operations"
    (assert-equal 1 (vm-eval '(first (list 1 2 3))))
    (assert-equal 3 (vm-eval '(len (list 1 2 3)))))

  (deftest "nested list"
    (assert-equal (list 1 (list 2 3))
      (vm-eval '(list 1 (list 2 3))))))


;; --------------------------------------------------------------------------
;; String operations
;; --------------------------------------------------------------------------

(defsuite "vm-strings"
  (deftest "str concat"
    (assert-equal "hello world" (vm-eval '(str "hello" " " "world"))))

  (deftest "string-length"
    (assert-equal 5 (vm-eval '(string-length "hello"))))

  (deftest "str coerces numbers"
    (assert-equal "42" (vm-eval '(str 42)))))


;; --------------------------------------------------------------------------
;; Define — top-level and local
;; --------------------------------------------------------------------------

(defsuite "vm-define"
  (deftest "top-level define"
    (assert-equal 42
      (vm-eval '(do (define x 42) x))))

  (deftest "define function then call"
    (assert-equal 10
      (vm-eval '(do
                  (define double (fn (n) (* n 2)))
                  (double 5)))))

  (deftest "local define inside fn"
    (assert-equal 30
      (vm-eval '(let ((f (fn (x)
                           (define y (* x 2))
                           (+ x y))))
                  (f 10)))))

  (deftest "define with forward reference"
    (assert-equal 120
      (vm-eval '(do
                  (define fact (fn (n)
                    (if (<= n 1) 1 (* n (fact (- n 1))))))
                  (fact 5))))))


;; --------------------------------------------------------------------------
;; Letrec — mutual recursion
;; --------------------------------------------------------------------------

(defsuite "vm-letrec"
  (deftest "letrec self-recursion"
    (assert-equal 55
      (vm-eval '(letrec ((sum-to (fn (n)
                                   (if (<= n 0) 0
                                     (+ n (sum-to (- n 1)))))))
                  (sum-to 10)))))

  (deftest "letrec mutual recursion"
    (assert-equal true
      (vm-eval '(letrec ((my-even? (fn (n)
                                     (if (= n 0) true (my-odd? (- n 1)))))
                          (my-odd? (fn (n)
                                    (if (= n 0) false (my-even? (- n 1))))))
                  (my-even? 10))))))


;; --------------------------------------------------------------------------
;; Quasiquote
;; --------------------------------------------------------------------------

(defsuite "vm-quasiquote"
  (deftest "simple quasiquote"
    (assert-equal (list 1 2 3)
      (vm-eval '(let ((x 2)) `(1 ,x 3)))))

  (deftest "quasiquote with splice"
    (assert-equal (list 1 2 3 4)
      (vm-eval '(let ((xs (list 2 3))) `(1 ,@xs 4))))))


;; --------------------------------------------------------------------------
;; Thread macro
;; --------------------------------------------------------------------------

(defsuite "vm-threading"
  (deftest "thread-first"
    (assert-equal 7
      (vm-eval '(-> 5 (+ 2)))))

  (deftest "thread-first chain"
    (assert-equal 12
      (vm-eval '(-> 10 (+ 5) (- 3))))))


;; --------------------------------------------------------------------------
;; Integration: compile then execute
;; --------------------------------------------------------------------------

(defsuite "vm-integration"
  (deftest "fibonacci"
    (assert-equal 55
      (vm-eval '(do
                  (define fib (fn (n)
                    (if (<= n 1) n
                      (+ (fib (- n 1)) (fib (- n 2))))))
                  (fib 10)))))

  (deftest "map via recursive define"
    (assert-equal (list 2 4 6)
      (vm-eval '(do
                  (define my-map (fn (f lst)
                    (if (empty? lst) (list)
                      (cons (f (first lst)) (my-map f (rest lst))))))
                  (my-map (fn (x) (* x 2)) (list 1 2 3))))))

  (deftest "filter via recursive define"
    (assert-equal (list 2 4)
      (vm-eval '(do
                  (define my-filter (fn (pred lst)
                    (if (empty? lst) (list)
                      (if (pred (first lst))
                        (cons (first lst) (my-filter pred (rest lst)))
                        (my-filter pred (rest lst))))))
                  (my-filter (fn (x) (even? x)) (list 1 2 3 4 5))))))

  (deftest "reduce via recursive define"
    (assert-equal 15
      (vm-eval '(do
                  (define my-reduce (fn (f acc lst)
                    (if (empty? lst) acc
                      (my-reduce f (f acc (first lst)) (rest lst)))))
                  (my-reduce (fn (acc x) (+ acc x)) 0 (list 1 2 3 4 5))))))

  (deftest "nested function calls"
    (assert-equal 42
      (vm-eval '(do
                  (define compose (fn (f g) (fn (x) (f (g x)))))
                  (define inc (fn (x) (+ x 1)))
                  (define double (fn (x) (* x 2)))
                  (let ((inc-then-double (compose double inc)))
                    (inc-then-double 20)))))))

;; --------------------------------------------------------------------------
;; VM recursive mutation — closure capture must preserve mutable references
;; --------------------------------------------------------------------------
;;
;; Regression: recursive functions that append! to a shared mutable list
;; lost mutations after the first call under JIT.  The stepper island's
;; split-tag function produced 1 step instead of 16, breaking SSR.

(defsuite "vm-recursive-mutation"
  (deftest "recursive append! to shared list"
    (assert-equal 3
      (vm-eval '(do
        (define walk (fn (items result)
          (when (not (empty? items))
            (append! result (first items))
            (walk (rest items) result))))
        (let ((result (list)))
          (walk (list "a" "b" "c") result)
          (len result))))))

  (deftest "recursive tree walk with append!"
    (assert-equal 7
      (vm-eval '(do
        (define walk-children (fn (items result walk-fn)
          (when (not (empty? items))
            (walk-fn (first items) result)
            (walk-children (rest items) result walk-fn))))
        (define walk (fn (expr result)
          (cond
            (not (list? expr))
              (append! result "leaf")
            (empty? expr) nil
            :else
              (do (append! result "open")
                  (walk-children (rest expr) result walk)
                  (append! result "close")))))
        (let ((tree (first (sx-parse "(div \"a\" (span \"b\") \"c\")")))
              (result (list)))
          (walk tree result)
          (len result))))))

  (deftest "recursive walk matching stepper split-tag pattern"
    (assert-equal 16
      (vm-eval '(do
        (define walk-each (fn (items result walk-fn)
          (when (not (empty? items))
            (walk-fn (first items) result)
            (walk-each (rest items) result walk-fn))))
        (define collect-children (fn (items cch)
          (when (not (empty? items))
            (let ((a (first items)))
              (if (and (list? a) (not (empty? a))
                       (= (type-of (first a)) "symbol")
                       (starts-with? (symbol-name (first a)) "~"))
                nil  ;; skip component spreads
                (append! cch a))
              (collect-children (rest items) cch)))))
        (define split-tag (fn (expr result)
          (cond
            (not (list? expr))
              (append! result "leaf")
            (empty? expr) nil
            (not (= (type-of (first expr)) "symbol"))
              (append! result "leaf")
            (is-html-tag? (symbol-name (first expr)))
              (let ((cch (list)))
                (collect-children (rest expr) cch)
                (append! result "open")
                (walk-each cch result split-tag)
                (append! result "close"))
            :else
              (append! result "expr"))))
        (let ((parsed (sx-parse "(div (~cssx/tw :tokens \"text-center\")\n  (h1 (~cssx/tw :tokens \"text-3xl font-bold mb-2\")\n    (span (~cssx/tw :tokens \"text-rose-500\") \"the \")\n    (span (~cssx/tw :tokens \"text-amber-500\") \"joy \")\n    (span (~cssx/tw :tokens \"text-emerald-500\") \"of \")\n    (span (~cssx/tw :tokens \"text-violet-600 text-4xl\") \"sx\")))"))
              (result (list)))
          (split-tag (first parsed) result)
          (len result)))))))