(define-library
  (sx vm)
  (export
    make-upvalue-cell
    uv-get
    uv-set!
    make-vm-code
    make-vm-closure
    make-vm-frame
    make-vm
    vm-push
    vm-pop
    vm-peek
    frame-read-u8
    frame-read-u16
    frame-read-i16
    vm-push-frame
    code-from-value
    vm-closure?
    *active-vm*
    *jit-compile-fn*
    lambda?
    lambda-compiled
    lambda-set-compiled!
    lambda-name
    cek-call-or-suspend
    try-jit-call
    vm-call
    frame-local-get
    frame-local-set
    frame-upvalue-get
    frame-upvalue-set
    frame-ip
    frame-set-ip!
    frame-base
    frame-closure
    closure-code
    closure-upvalues
    closure-env
    code-bytecode
    code-constants
    code-locals
    vm-sp
    vm-set-sp!
    vm-stack
    vm-set-stack!
    vm-frames
    vm-set-frames!
    vm-globals-ref
    collect-n-from-stack
    collect-n-pairs
    pad-n-nils
    vm-global-get
    vm-resolve-ho-form
    vm-call-external
    vm-global-set
    env-walk
    env-walk-set!
    vm-create-closure
    vm-run
    vm-step
    vm-call-closure
    vm-execute-module
    vm-resume-module)
  (begin
    (define make-upvalue-cell (fn (value) {:uv-value value}))
    (define uv-get (fn (cell) (get cell "uv-value")))
    (define uv-set! (fn (cell value) (dict-set! cell "uv-value" value)))
    (define make-vm-code (fn (arity locals bytecode constants) {:vc-bytecode bytecode :vc-locals locals :vc-arity arity :vc-constants constants}))
    (define
      make-vm-closure
      (fn (code upvalues name globals closure-env) {:vm-globals globals :vm-upvalues upvalues :vm-name name :vm-code code :vm-closure-env closure-env}))
    (define make-vm-frame (fn (closure base) {:ip 0 :closure closure :base base :local-cells {}}))
    (define make-vm (fn (globals) {:sp 0 :frames (list) :stack (make-vm-stack 4096) :globals globals}))
    (define
      vm-push
      (fn
        (vm value)
        (let
          ((sp (vm-sp vm)) (stack (vm-stack vm)))
          (when
            (>= sp (vm-stack-length stack))
            (let
              ((new-stack (vm-stack-grow stack sp)))
              (vm-stack-copy! stack new-stack sp)
              (vm-set-stack! vm new-stack)
              (set! stack new-stack)))
          (vm-stack-set! stack sp value)
          (vm-set-sp! vm (+ sp 1)))))
    (define
      vm-pop
      (fn
        (vm)
        (let
          ((sp (- (vm-sp vm) 1)))
          (vm-set-sp! vm sp)
          (vm-stack-get (vm-stack vm) sp))))
    (define
      vm-peek
      (fn (vm) (vm-stack-get (vm-stack vm) (- (vm-sp vm) 1))))
    (define
      frame-read-u8
      (fn
        (frame)
        (let
          ((ip (frame-ip frame))
            (bc (-> frame frame-closure closure-code code-bytecode)))
          (let ((v (nth bc ip))) (frame-set-ip! frame (+ ip 1)) v))))
    (define
      frame-read-u16
      (fn
        (frame)
        (let
          ((lo (frame-read-u8 frame)) (hi (frame-read-u8 frame)))
          (+ lo (* hi 256)))))
    (define
      frame-read-i16
      (fn
        (frame)
        (let
          ((v (frame-read-u16 frame)))
          (if (>= v 32768) (- v 65536) v))))
    (define
      vm-push-frame
      (fn
        (vm closure args)
        (let
          ((frame (make-vm-frame closure (vm-sp vm))))
          (for-each (fn (a) (vm-push vm a)) args)
          (pad-n-nils
            vm
            (- (code-locals (closure-code closure)) (len args)))
          (vm-set-frames! vm (cons frame (vm-frames vm))))))
    (define
      code-from-value
      (fn
        (v)
        "Convert a compiler output dict to a vm-code object."
        (if
          (not (dict? v))
          (make-vm-code 0 16 (list) (list))
          (let
            ((bc-raw (get v "bytecode"))
              (bc (if (nil? bc-raw) (list) bc-raw))
              (consts-raw (get v "constants"))
              (consts (if (nil? consts-raw) (list) consts-raw))
              (arity-raw (get v "arity"))
              (arity (if (nil? arity-raw) 0 arity-raw)))
            (make-vm-code arity (+ arity 16) bc consts)))))
    (define vm-closure? (fn (v) (and (dict? v) (has-key? v "vm-code"))))
    (define *active-vm* nil)
    (define *jit-compile-fn* nil)
    (define lambda? (fn (v) (= (type-of v) "lambda")))
    (define lambda-compiled (fn (f) nil))
    (define lambda-set-compiled! (fn (f val) nil))
    (define lambda-name (fn (f) nil))
    (define cek-call-or-suspend (fn (vm f args) (cek-call f args)))
    (define
      try-jit-call
      (fn
        (vm f args)
        (let
          ((compiled (lambda-compiled f)))
          (cond
            (vm-closure? compiled)
            (vm-push vm (vm-call-closure compiled args (vm-globals-ref vm)))
            (= compiled :jit-failed)
            (vm-push vm (cek-call-or-suspend vm f args))
            (and *jit-compile-fn* (lambda-name f))
            (do
              (lambda-set-compiled! f :jit-failed)
              (let
                ((result (*jit-compile-fn* f (vm-globals-ref vm))))
                (if
                  (vm-closure? result)
                  (do
                    (lambda-set-compiled! f result)
                    (vm-push
                      vm
                      (vm-call-closure result args (vm-globals-ref vm))))
                  (vm-push vm (cek-call-or-suspend vm f args)))))
            :else (vm-push vm (cek-call-or-suspend vm f args))))))
    (define
      vm-call
      (fn
        (vm f args)
        (cond
          (vm-closure? f)
          (vm-push-frame vm f args)
          (lambda? f)
          (try-jit-call vm f args)
          (or (= (type-of f) "component") (= (type-of f) "island"))
          (vm-push vm (cek-call-or-suspend vm f args))
          (callable? f)
          (vm-push vm (apply f args))
          :else (error (str "VM: not callable: " (type-of f))))))
    (define
      frame-local-get
      (fn
        (vm frame slot)
        "Read a local variable — check shared cells first, then stack."
        (let
          ((cells (get frame "local-cells")) (key (str slot)))
          (if
            (has-key? cells key)
            (uv-get (get cells key))
            (vm-stack-get (vm-stack vm) (+ (frame-base frame) slot))))))
    (define
      frame-local-set
      (fn
        (vm frame slot value)
        "Write a local variable — to shared cell or stack."
        (let
          ((cells (get frame "local-cells")) (key (str slot)))
          (if
            (has-key? cells key)
            (uv-set! (get cells key) value)
            (vm-stack-set! (vm-stack vm) (+ (frame-base frame) slot) value)))))
    (define
      frame-upvalue-get
      (fn
        (frame idx)
        (uv-get (nth (-> frame frame-closure closure-upvalues) idx))))
    (define
      frame-upvalue-set
      (fn
        (frame idx value)
        (uv-set! (nth (-> frame frame-closure closure-upvalues) idx) value)))
    (define frame-ip (fn (frame) (get frame "ip")))
    (define frame-set-ip! (fn (frame val) (dict-set! frame "ip" val)))
    (define frame-base (fn (frame) (get frame "base")))
    (define frame-closure (fn (frame) (get frame "closure")))
    (define closure-code (fn (cl) (get cl "vm-code")))
    (define closure-upvalues (fn (cl) (get cl "vm-upvalues")))
    (define closure-env (fn (cl) (get cl "closure-env")))
    (define code-bytecode (fn (code) (get code "vc-bytecode")))
    (define code-constants (fn (code) (get code "vc-constants")))
    (define code-locals (fn (code) (get code "vc-locals")))
    (define vm-sp (fn (vm) (get vm "sp")))
    (define vm-set-sp! (fn (vm val) (dict-set! vm "sp" val)))
    (define vm-stack (fn (vm) (get vm "stack")))
    (define vm-set-stack! (fn (vm val) (dict-set! vm "stack" val)))
    (define vm-frames (fn (vm) (get vm "frames")))
    (define vm-set-frames! (fn (vm val) (dict-set! vm "frames" val)))
    (define vm-globals-ref (fn (vm) (get vm "globals")))
    (define
      collect-n-from-stack
      (fn
        (vm n)
        (let
          ((result (list)) (i 0))
          (define
            _loop
            (fn
              ()
              (when
                (< i n)
                (set! result (cons (vm-pop vm) result))
                (set! i (+ i 1))
                (_loop))))
          (_loop)
          result)))
    (define
      pad-n-nils
      (fn
        (vm n)
        (let
          ((i 0))
          (define
            _loop
            (fn
              ()
              (when (< i n) (vm-push vm nil) (set! i (+ i 1)) (_loop))))
          (_loop))))
    (define
      collect-n-pairs
      (fn
        (vm n)
        (let
          ((d {}) (i 0))
          (define
            _loop
            (fn
              ()
              (when
                (< i n)
                (let
                  ((v (vm-pop vm)) (k (vm-pop vm)))
                  (dict-set! d k v)
                  (set! i (+ i 1))
                  (_loop)))))
          (_loop)
          d)))
    (define
      vm-global-get
      (fn
        (vm frame name)
        "Look up a global: globals table → closure env → primitives → HO wrappers"
        (let
          ((globals (vm-globals-ref vm)))
          (if
            (has-key? globals name)
            (get globals name)
            (let
              ((closure-env (-> frame frame-closure closure-env)))
              (if
                (nil? closure-env)
                (cek-try
                  (fn () (get-primitive name))
                  (fn (e) (vm-resolve-ho-form vm name)))
                (let
                  ((found (env-walk closure-env name)))
                  (if
                    (nil? found)
                    (cek-try
                      (fn () (get-primitive name))
                      (fn (e) (vm-resolve-ho-form vm name)))
                    found))))))))
    (define
      vm-resolve-ho-form
      (fn
        (vm name)
        (match
          name
          ("for-each"
            (fn
              (f coll)
              (for-each (fn (x) (vm-call-external vm f (list x))) coll)))
          ("map"
            (fn
              (f coll)
              (map (fn (x) (vm-call-external vm f (list x))) coll)))
          ("map-indexed"
            (fn
              (f coll)
              (map-indexed
                (fn (i x) (vm-call-external vm f (list i x)))
                coll)))
          ("filter"
            (fn
              (f coll)
              (filter (fn (x) (vm-call-external vm f (list x))) coll)))
          ("reduce"
            (fn
              (f init coll)
              (reduce
                (fn (acc x) (vm-call-external vm f (list acc x)))
                init
                coll)))
          ("some"
            (fn
              (f coll)
              (some (fn (x) (vm-call-external vm f (list x))) coll)))
          ("every?"
            (fn
              (f coll)
              (every? (fn (x) (vm-call-external vm f (list x))) coll)))
          (_ (error (str "VM undefined: " name))))))
    (define
      vm-call-external
      (fn
        (vm f args)
        (if
          (vm-closure? f)
          (vm-call-closure f args (vm-globals-ref vm))
          (cek-call f args))))
    (define
      vm-global-set
      (fn
        (vm frame name value)
        "Set a global: write to closure env if found, else globals table."
        (let
          ((closure-env (get (frame-closure frame) "vm-closure-env"))
            (written false))
          (when
            (not (nil? closure-env))
            (set! written (env-walk-set! closure-env name value)))
          (when (not written) (dict-set! (vm-globals-ref vm) name value)))))
    (define
      env-walk
      (fn
        (env name)
        (if
          (nil? env)
          nil
          (if
            (env-has? env name)
            (env-get env name)
            (let
              ((parent (env-parent env)))
              (if (nil? parent) nil (env-walk parent name)))))))
    (define
      env-walk-set!
      (fn
        (env name value)
        (if
          (nil? env)
          false
          (if
            (env-has? env name)
            (do (env-set! env name value) true)
            (let
              ((parent (env-parent env)))
              (if (nil? parent) false (env-walk-set! parent name value)))))))
    (define
      vm-create-closure
      (fn
        (vm frame code-val)
        "Create a closure from a code constant. Reads upvalue descriptors\n     from the bytecode stream and captures values from the enclosing frame."
        (let
          ((code (code-from-value code-val))
            (uv-count
              (if (dict? code-val) (or (get code-val "upvalue-count") 0) 0)))
          (let
            ((upvalues (let ((result (list)) (i 0)) (define capture-loop (fn () (when (< i uv-count) (let ((is-local (frame-read-u8 frame)) (index (frame-read-u8 frame))) (let ((cell (if (= is-local 1) (let ((cells (get frame "local-cells")) (key (str index))) (if (has-key? cells key) (get cells key) (let ((c (make-upvalue-cell (vm-stack-get (vm-stack vm) (+ (frame-base frame) index))))) (dict-set! cells key c) c))) (nth (-> frame frame-closure closure-upvalues) index)))) (append! result cell) (set! i (+ i 1)) (capture-loop)))))) (capture-loop) result)))
            (make-vm-closure code upvalues nil (vm-globals-ref vm) nil)))))
    (define
      vm-run
      (fn
        (vm)
        "Execute bytecode until all frames are done or IO suspension."
        (define
          loop
          (fn
            ()
            (when
              (not (empty? (vm-frames vm)))
              (let
                ((frame (first (vm-frames vm)))
                  (rest-frames (rest (vm-frames vm))))
                (let
                  ((bc (-> frame frame-closure closure-code code-bytecode))
                    (consts
                      (-> frame frame-closure closure-code code-constants)))
                  (if
                    (>= (frame-ip frame) (len bc))
                    (vm-set-frames! vm (list))
                    (do
                      (vm-step vm frame rest-frames bc consts)
                      (when
                        (nil? (get (vm-globals-ref vm) "__io_request"))
                        (loop)))))))))
        (loop)))
    (define
      vm-step
      (fn
        (vm frame rest-frames bc consts)
        (let
          ((op (frame-read-u8 frame)))
          (cond
            (= op 1)
            (let
              ((idx (frame-read-u16 frame)))
              (vm-push vm (nth consts idx)))
            (= op 2)
            (vm-push vm nil)
            (= op 3)
            (vm-push vm true)
            (= op 4)
            (vm-push vm false)
            (= op 5)
            (vm-pop vm)
            (= op 6)
            (vm-push vm (vm-peek vm))
            (= op 16)
            (let
              ((slot (frame-read-u8 frame)))
              (vm-push vm (frame-local-get vm frame slot)))
            (= op 17)
            (let
              ((slot (frame-read-u8 frame)))
              (frame-local-set vm frame slot (vm-peek vm)))
            (= op 18)
            (let
              ((idx (frame-read-u8 frame)))
              (vm-push vm (frame-upvalue-get frame idx)))
            (= op 19)
            (let
              ((idx (frame-read-u8 frame)))
              (frame-upvalue-set frame idx (vm-peek vm)))
            (= op 20)
            (let
              ((idx (frame-read-u16 frame)) (name (nth consts idx)))
              (vm-push vm (vm-global-get vm frame name)))
            (= op 21)
            (let
              ((idx (frame-read-u16 frame)) (name (nth consts idx)))
              (vm-global-set vm frame name (vm-peek vm)))
            (= op 32)
            (let
              ((offset (frame-read-i16 frame)))
              (frame-set-ip! frame (+ (frame-ip frame) offset)))
            (= op 33)
            (let
              ((offset (frame-read-i16 frame)) (v (vm-pop vm)))
              (when
                (not v)
                (frame-set-ip! frame (+ (frame-ip frame) offset))))
            (= op 34)
            (let
              ((offset (frame-read-i16 frame)) (v (vm-pop vm)))
              (when v (frame-set-ip! frame (+ (frame-ip frame) offset))))
            (= op 48)
            (let
              ((argc (frame-read-u8 frame))
                (args (collect-n-from-stack vm argc))
                (f (vm-pop vm)))
              (vm-call vm f args))
            (= op 49)
            (let
              ((argc (frame-read-u8 frame))
                (args (collect-n-from-stack vm argc))
                (f (vm-pop vm)))
              (vm-set-frames! vm rest-frames)
              (vm-set-sp! vm (frame-base frame))
              (vm-call vm f args))
            (= op 50)
            (let
              ((result (vm-pop vm)))
              (vm-set-frames! vm rest-frames)
              (vm-set-sp! vm (frame-base frame))
              (vm-push vm result))
            (= op 51)
            (let
              ((idx (frame-read-u16 frame)) (code-val (nth consts idx)))
              (let
                ((cl (vm-create-closure vm frame code-val)))
                (vm-push vm cl)))
            (= op 52)
            (let
              ((idx (frame-read-u16 frame))
                (argc (frame-read-u8 frame))
                (name (nth consts idx))
                (args (collect-n-from-stack vm argc)))
              (vm-push vm (call-primitive name args)))
            (= op 64)
            (let
              ((count (frame-read-u16 frame))
                (items (collect-n-from-stack vm count)))
              (vm-push vm items))
            (= op 65)
            (let
              ((count (frame-read-u16 frame))
                (d (collect-n-pairs vm count)))
              (vm-push vm d))
            (= op 144)
            (let
              ((count (frame-read-u8 frame))
                (parts (collect-n-from-stack vm count)))
              (vm-push vm (apply str parts)))
            (= op 128)
            (let
              ((idx (frame-read-u16 frame)) (name (nth consts idx)))
              (dict-set! (vm-globals-ref vm) name (vm-peek vm)))
            (= op 160)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (+ a b)))
            (= op 161)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (- a b)))
            (= op 162)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (* a b)))
            (= op 163)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (/ a b)))
            (= op 164)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (= a b)))
            (= op 165)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (< a b)))
            (= op 166)
            (let ((b (vm-pop vm)) (a (vm-pop vm))) (vm-push vm (> a b)))
            (= op 167)
            (vm-push vm (not (vm-pop vm)))
            (= op 168)
            (vm-push vm (len (vm-pop vm)))
            (= op 169)
            (vm-push vm (first (vm-pop vm)))
            (= op 170)
            (vm-push vm (rest (vm-pop vm)))
            (= op 171)
            (let
              ((n (vm-pop vm)) (coll (vm-pop vm)))
              (vm-push vm (nth coll n)))
            (= op 172)
            (let
              ((coll (vm-pop vm)) (x (vm-pop vm)))
              (vm-push vm (cons x coll)))
            (= op 173)
            (vm-push vm (- 0 (vm-pop vm)))
            (= op 174)
            (vm-push vm (inc (vm-pop vm)))
            (= op 175)
            (vm-push vm (dec (vm-pop vm)))
            (= op 112)
            (let
              ((request (vm-pop vm)))
              (dict-set! (vm-globals-ref vm) "__io_request" request))
            :else (error (str "VM: unknown opcode " op))))))
    (define
      vm-call-closure
      (fn
        (closure args globals)
        (let
          ((prev-vm *active-vm*) (vm (make-vm globals)))
          (set! *active-vm* vm)
          (vm-push-frame vm closure args)
          (vm-run vm)
          (set! *active-vm* prev-vm)
          (vm-pop vm))))
    (define
      vm-execute-module
      (fn
        (code globals)
        (let
          ((vm-code (code-from-value code)) (vm (make-vm globals)))
          (let
            ((closure (make-vm-closure vm-code (list) "module" globals nil))
              (frame (make-vm-frame closure 0)))
            (pad-n-nils vm (code-locals vm-code))
            (vm-set-frames! vm (list frame))
            (vm-run vm)
            (let
              ((io-req (get (vm-globals-ref vm) "__io_request")))
              (if (nil? io-req) (vm-pop vm) {:vm vm :suspended true :op "import" :request io-req}))))))
    (define
      vm-resume-module
      (fn
        (suspended-result)
        "Resume a suspended VM after IO (import) has been resolved.\n     Clears __io_request in globals, pushes nil (import result), re-runs."
        (let
          ((vm (get suspended-result :vm)))
          (dict-set! (vm-globals-ref vm) "__io_request" nil)
          (vm-push vm nil)
          (vm-run vm)
          (let
            ((io-req (get (vm-globals-ref vm) "__io_request")))
            (if (nil? io-req) (vm-pop vm) {:vm vm :suspended true :op "import" :request io-req}))))))) ;; end define-library

;; Re-export to global namespace for backward compatibility
(import (sx vm))