(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?
    vm-call
    frame-local-get
    frame-local-set
    frame-upvalue-get
    frame-upvalue-set
    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)
  (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 (get vm "sp")) (stack (get vm "stack")))
      (when
        (>= sp (vm-stack-length stack))
        (let
          ((new-stack (make-vm-stack (* sp 2))))
          (vm-stack-copy! stack new-stack sp)
          (dict-set! vm "stack" new-stack)
          (set! stack new-stack)))
      (vm-stack-set! stack sp value)
      (dict-set! vm "sp" (+ sp 1)))))

(define
  vm-pop
  (fn
    (vm)
    (let
      ((sp (- (get vm "sp") 1)))
      (dict-set! vm "sp" sp)
      (vm-stack-get (get vm "stack") sp))))

(define
  vm-peek
  (fn (vm) (vm-stack-get (get vm "stack") (- (get vm "sp") 1))))

(define
  frame-read-u8
  (fn
    (frame)
    (let
      ((ip (get frame "ip"))
        (bc (get (get (get frame "closure") "vm-code") "vc-bytecode")))
      (let ((v (nth bc ip))) (dict-set! frame "ip" (+ 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 (get vm "sp"))))
      (for-each (fn (a) (vm-push vm a)) args)
      (let
        ((arity (len args))
          (total-locals (get (get closure "vm-code") "vc-locals")))
        (let
          ((pad-count (- total-locals arity)))
          (when
            (> pad-count 0)
            (let
              ((i 0))
              (define
                pad-loop
                (fn
                  ()
                  (when
                    (< i pad-count)
                    (vm-push vm nil)
                    (set! i (+ i 1))
                    (pad-loop))))
              (pad-loop)))))
      (dict-set! vm "frames" (cons frame (get vm "frames"))))))

(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
  vm-call
  (fn
    (vm f args)
    (cond
      (vm-closure? f)
      (vm-push-frame vm f args)
      (or
        (= (type-of f) "lambda")
        (= (type-of f) "component")
        (= (type-of f) "island"))
      (vm-push vm (cek-call 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 (get vm "stack") (+ (get frame "base") slot))))))

(define
  frame-local-set
  (fn
    (vm frame slot value)
    "Write a local variable — to shared cell if captured, else to stack."
    (let
      ((cells (get frame "local-cells")) (key (str slot)))
      (if
        (has-key? cells key)
        (uv-set! (get cells key) value)
        (vm-stack-set! (get vm "stack") (+ (get frame "base") slot) value)))))

(define
  frame-upvalue-get
  (fn
    (frame idx)
    (uv-get (nth (get (get frame "closure") "vm-upvalues") idx))))

(define
  frame-upvalue-set
  (fn
    (frame idx value)
    (uv-set! (nth (get (get frame "closure") "vm-upvalues") idx) value)))

(define
  vm-global-get
  (fn
    (vm frame name)
    "Look up a global: globals table → closure env → primitives → HO wrappers"
    (let
      ((globals (get vm "globals")))
      (if
        (has-key? globals name)
        (get globals name)
        (let
          ((closure-env (get (get 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)
    (cond
      (= name "for-each")
      (fn
        (f coll)
        (for-each (fn (x) (vm-call-external vm f (list x))) coll))
      (= name "map")
      (fn (f coll) (map (fn (x) (vm-call-external vm f (list x))) coll))
      (= name "map-indexed")
      (fn
        (f coll)
        (map-indexed (fn (i x) (vm-call-external vm f (list i x))) coll))
      (= name "filter")
      (fn
        (f coll)
        (filter (fn (x) (vm-call-external vm f (list x))) coll))
      (= name "reduce")
      (fn
        (f init coll)
        (reduce
          (fn (acc x) (vm-call-external vm f (list acc x)))
          init
          coll))
      (= name "some")
      (fn (f coll) (some (fn (x) (vm-call-external vm f (list x))) coll))
      (= name "every?")
      (fn
        (f coll)
        (every? (fn (x) (vm-call-external vm f (list x))) coll))
      :else (error (str "VM undefined: " name)))))

(define
  vm-call-external
  (fn
    (vm f args)
    (if
      (vm-closure? f)
      (vm-call-closure f args (get vm "globals"))
      (cek-call f args))))

(define
  vm-global-set
  (fn
    (vm frame name value)
    "Set a global: write to closure env if name exists there, else globals."
    (let
      ((closure-env (get (get frame "closure") "vm-closure-env"))
        (written false))
      (when
        (not (nil? closure-env))
        (set! written (env-walk-set! closure-env name value)))
      (when (not written) (dict-set! (get vm "globals") 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)
            (let ((n (get code-val "upvalue-count"))) (if (nil? n) 0 n))
            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 (get vm "stack") (+ (get frame "base") index))))) (dict-set! cells key c) c))) (nth (get (get frame "closure") "vm-upvalues") index)))) (append! result cell) (set! i (+ i 1)) (capture-loop)))))) (capture-loop) result)))
        (make-vm-closure code upvalues nil (get vm "globals") nil)))))

(define
  vm-run
  (fn
    (vm)
    "Execute bytecode until all frames are exhausted.\n     VmClosure calls push new frames; the loop picks them up.\n     OP_TAIL_CALL + VmClosure = true TCO: drop frame, push new, loop."
    (define
      loop
      (fn
        ()
        (when
          (not (empty? (get vm "frames")))
          (let
            ((frame (first (get vm "frames")))
              (rest-frames (rest (get vm "frames"))))
            (let
              ((bc (get (get (get frame "closure") "vm-code") "vc-bytecode"))
                (consts
                  (get (get (get frame "closure") "vm-code") "vc-constants")))
              (if
                (>= (get frame "ip") (len bc))
                (dict-set! vm "frames" (list))
                (do (vm-step vm frame rest-frames bc consts) (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)))
          (dict-set! frame "ip" (+ (get frame "ip") offset)))
        (= op 33)
        (let
          ((offset (frame-read-i16 frame)) (v (vm-pop vm)))
          (when (not v) (dict-set! frame "ip" (+ (get frame "ip") offset))))
        (= op 34)
        (let
          ((offset (frame-read-i16 frame)) (v (vm-pop vm)))
          (when v (dict-set! frame "ip" (+ (get frame "ip") offset))))
        (= op 48)
        (let
          ((argc (frame-read-u8 frame)) (args-rev (list)) (i 0))
          (define
            collect-args
            (fn
              ()
              (when
                (< i argc)
                (set! args-rev (cons (vm-pop vm) args-rev))
                (set! i (+ i 1))
                (collect-args))))
          (collect-args)
          (let ((f (vm-pop vm))) (vm-call vm f args-rev)))
        (= op 49)
        (let
          ((argc (frame-read-u8 frame)) (args-rev (list)) (i 0))
          (define
            collect-args
            (fn
              ()
              (when
                (< i argc)
                (set! args-rev (cons (vm-pop vm) args-rev))
                (set! i (+ i 1))
                (collect-args))))
          (collect-args)
          (let
            ((f (vm-pop vm)))
            (dict-set! vm "frames" rest-frames)
            (dict-set! vm "sp" (get frame "base"))
            (vm-call vm f args-rev)))
        (= op 50)
        (let
          ((result (vm-pop vm)))
          (dict-set! vm "frames" rest-frames)
          (dict-set! vm "sp" (get frame "base"))
          (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-rev (list))
            (i 0))
          (define
            collect-args
            (fn
              ()
              (when
                (< i argc)
                (set! args-rev (cons (vm-pop vm) args-rev))
                (set! i (+ i 1))
                (collect-args))))
          (collect-args)
          (vm-push vm (call-primitive name args-rev)))
        (= op 64)
        (let
          ((count (frame-read-u16 frame)) (items-rev (list)) (i 0))
          (define
            collect-items
            (fn
              ()
              (when
                (< i count)
                (set! items-rev (cons (vm-pop vm) items-rev))
                (set! i (+ i 1))
                (collect-items))))
          (collect-items)
          (vm-push vm items-rev))
        (= op 65)
        (let
          ((count (frame-read-u16 frame)) (d {}) (i 0))
          (define
            collect-pairs
            (fn
              ()
              (when
                (< i count)
                (let
                  ((v (vm-pop vm)) (k (vm-pop vm)))
                  (dict-set! d k v)
                  (set! i (+ i 1))
                  (collect-pairs)))))
          (collect-pairs)
          (vm-push vm d))
        (= op 144)
        (let
          ((count (frame-read-u8 frame)) (parts-rev (list)) (i 0))
          (define
            collect-parts
            (fn
              ()
              (when
                (< i count)
                (set! parts-rev (cons (vm-pop vm) parts-rev))
                (set! i (+ i 1))
                (collect-parts))))
          (collect-parts)
          (vm-push vm (apply str parts-rev)))
        (= op 128)
        (let
          ((idx (frame-read-u16 frame)) (name (nth consts idx)))
          (dict-set! (get vm "globals") 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)))
          (error (str "VM: IO suspension (OP_PERFORM) — request: " request)))
        :else (error (str "VM: unknown opcode " op))))))

(define
  vm-call-closure
  (fn
    (closure args globals)
    (let
      ((vm (make-vm globals)))
      (vm-push-frame vm closure args)
      (vm-run vm)
      (vm-pop vm))))

(define
  vm-execute-module
  (fn
    (code globals)
    (let
      ((closure (make-vm-closure code (list) "module" globals nil))
        (vm (make-vm globals)))
      (let
        ((frame (make-vm-frame closure 0)))
        (let
          ((i 0) (total (get code "vc-locals")))
          (define
            pad-loop
            (fn
              ()
              (when
                (< i total)
                (vm-push vm nil)
                (set! i (+ i 1))
                (pad-loop))))
          (pad-loop))
        (dict-set! vm "frames" (list frame))
        (vm-run vm)
        (vm-pop vm)))))


)) ;; end define-library

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