;; ==========================================================================
;; vm.sx — SX bytecode virtual machine
;;
;; Stack-based interpreter for bytecode produced by compiler.sx.
;; Written in SX — transpiled to each target (OCaml, JS, WASM).
;;
;; Architecture:
;;   - Array-based value stack (no allocation per step)
;;   - Frame list for call stack (one frame per function invocation)
;;   - Upvalue cells for shared mutable closure variables
;;   - Iterative dispatch loop (no host-stack growth)
;;   - TCO via frame replacement on OP_TAIL_CALL
;;
;; Platform interface:
;;   The host must provide:
;;     - make-vm-stack, vm-stack-get, vm-stack-set!, vm-stack-grow
;;     - cek-call (fallback for Lambda/Component)
;;     - get-primitive (primitive lookup)
;;   Everything else is defined here.
;; ==========================================================================


;; --------------------------------------------------------------------------
;; 1. Types — VM data structures
;; --------------------------------------------------------------------------

;; Upvalue cell — shared mutable reference for captured variables.
;; When a closure captures a local, both the parent frame and the
;; closure read/write through this cell.
(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)))

;; VM code object — compiled bytecode + constant pool.
;; Produced by compiler.sx, consumed by the VM.
(define make-vm-code
  (fn (arity locals bytecode constants)
    {:vc-arity arity
     :vc-locals locals
     :vc-bytecode bytecode
     :vc-constants constants}))

;; VM closure — code + captured upvalues + globals reference.
(define make-vm-closure
  (fn (code upvalues name globals closure-env)
    {:vm-code code
     :vm-upvalues upvalues
     :vm-name name
     :vm-globals globals
     :vm-closure-env closure-env}))

;; VM frame — one per active function invocation.
(define make-vm-frame
  (fn (closure base)
    {:closure closure
     :ip 0
     :base base
     :local-cells {}}))

;; VM state — the virtual machine.
(define make-vm
  (fn (globals)
    {:stack (make-vm-stack 4096)
     :sp 0
     :frames (list)
     :globals globals}))


;; --------------------------------------------------------------------------
;; 2. Stack operations
;; --------------------------------------------------------------------------

(define vm-push
  (fn (vm value)
    (let ((sp (get vm "sp"))
          (stack (get vm "stack")))
      ;; Grow stack if needed
      (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))))


;; --------------------------------------------------------------------------
;; 3. Operand reading — read from bytecode stream
;; --------------------------------------------------------------------------

(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))))


;; --------------------------------------------------------------------------
;; 4. Frame management
;; --------------------------------------------------------------------------

;; Push a closure frame onto the VM.
;; Lays out args as locals, pads remaining locals with nil.
(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)
      ;; Pad remaining local slots with nil
      (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"))))))


;; --------------------------------------------------------------------------
;; 5. Code loading — convert compiler output to VM structures
;; --------------------------------------------------------------------------

(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)))))


;; --------------------------------------------------------------------------
;; 6. Call dispatch — route calls by value type
;; --------------------------------------------------------------------------

;; vm-call dispatches a function call within the VM.
;; VmClosure: push frame on current VM (fast path, enables TCO).
;; NativeFn: call directly, push result.
;; Lambda/Component: fall back to CEK evaluator.
(define vm-closure?
  (fn (v)
    (and (dict? v) (has-key? v "vm-code"))))

(define vm-call
  (fn (vm f args)
    (cond
      (vm-closure? f)
        ;; Fast path: push frame on current VM
        (vm-push-frame vm f args)

      (callable? f)
        ;; Native function or primitive
        (vm-push vm (apply f args))

      (or (= (type-of f) "lambda") (= (type-of f) "component") (= (type-of f) "island"))
        ;; CEK fallback — the host provides cek-call
        (vm-push vm (cek-call f args))

      :else
        (error (str "VM: not callable: " (type-of f))))))


;; --------------------------------------------------------------------------
;; 7. Local/upvalue access helpers
;; --------------------------------------------------------------------------

(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)))


;; --------------------------------------------------------------------------
;; 8. Global variable access with closure env chain
;; --------------------------------------------------------------------------

(define vm-global-get
  (fn (vm frame name)
    "Look up a global: globals table → closure env chain → primitives."
    (let ((globals (get vm "globals")))
      (if (has-key? globals name)
        (get globals name)
        ;; Walk the closure env chain for inner functions
        (let ((closure-env (get (get frame "closure") "vm-closure-env")))
          (if (nil? closure-env)
            (get-primitive name)
            (let ((found (env-walk closure-env name)))
              (if (nil? found)
                (get-primitive name)
                found))))))))

(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)))))

;; env-walk: walk an environment chain looking for a binding.
;; Returns the value or nil if not found.
(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)))))))

;; env-walk-set!: walk an environment chain, set value if name found.
;; Returns true if set, false if not found.
(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)))))))


;; --------------------------------------------------------------------------
;; 9. Closure creation — OP_CLOSURE with upvalue capture
;; --------------------------------------------------------------------------

(define vm-create-closure
  (fn (vm frame code-val)
    "Create a closure from a code constant. Reads upvalue descriptors
     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)
                                  ;; Capture from enclosing frame's local slot.
                                  ;; Create/reuse a shared cell so both parent
                                  ;; and closure read/write through it.
                                  (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)))
                                  ;; Capture from enclosing frame's upvalue
                                  (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)))))


;; --------------------------------------------------------------------------
;; 10. Main execution loop — iterative dispatch
;; --------------------------------------------------------------------------

(define vm-run
  (fn (vm)
    "Execute bytecode until all frames are exhausted.
     VmClosure calls push new frames; the loop picks them up.
     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))
                ;; Bytecode exhausted — stop
                (dict-set! vm "frames" (list))
                (do
                  (vm-step vm frame rest-frames bc consts)
                  (loop))))))))
    (loop)))


;; --------------------------------------------------------------------------
;; 11. Single step — opcode dispatch
;; --------------------------------------------------------------------------

(define vm-step
  (fn (vm frame rest-frames bc consts)
    (let ((op (frame-read-u8 frame)))
      (cond

        ;; ---- Constants ----
        (= op 1)   ;; OP_CONST
          (let ((idx (frame-read-u16 frame)))
            (vm-push vm (nth consts idx)))

        (= op 2)   ;; OP_NIL
          (vm-push vm nil)

        (= op 3)   ;; OP_TRUE
          (vm-push vm true)

        (= op 4)   ;; OP_FALSE
          (vm-push vm false)

        (= op 5)   ;; OP_POP
          (vm-pop vm)

        (= op 6)   ;; OP_DUP
          (vm-push vm (vm-peek vm))

        ;; ---- Variable access ----
        (= op 16)  ;; OP_LOCAL_GET
          (let ((slot (frame-read-u8 frame)))
            (vm-push vm (frame-local-get vm frame slot)))

        (= op 17)  ;; OP_LOCAL_SET
          (let ((slot (frame-read-u8 frame)))
            (frame-local-set vm frame slot (vm-peek vm)))

        (= op 18)  ;; OP_UPVALUE_GET
          (let ((idx (frame-read-u8 frame)))
            (vm-push vm (frame-upvalue-get frame idx)))

        (= op 19)  ;; OP_UPVALUE_SET
          (let ((idx (frame-read-u8 frame)))
            (frame-upvalue-set frame idx (vm-peek vm)))

        (= op 20)  ;; OP_GLOBAL_GET
          (let ((idx (frame-read-u16 frame))
                (name (nth consts idx)))
            (vm-push vm (vm-global-get vm frame name)))

        (= op 21)  ;; OP_GLOBAL_SET
          (let ((idx (frame-read-u16 frame))
                (name (nth consts idx)))
            (vm-global-set vm frame name (vm-peek vm)))

        ;; ---- Control flow ----
        (= op 32)  ;; OP_JUMP
          (let ((offset (frame-read-i16 frame)))
            (dict-set! frame "ip" (+ (get frame "ip") offset)))

        (= op 33)  ;; OP_JUMP_IF_FALSE
          (let ((offset (frame-read-i16 frame))
                (v (vm-pop vm)))
            (when (not v)
              (dict-set! frame "ip" (+ (get frame "ip") offset))))

        (= op 34)  ;; OP_JUMP_IF_TRUE
          (let ((offset (frame-read-i16 frame))
                (v (vm-pop vm)))
            (when v
              (dict-set! frame "ip" (+ (get frame "ip") offset))))

        ;; ---- Function calls ----
        (= op 48)  ;; OP_CALL
          (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)  ;; OP_TAIL_CALL
          (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)))
              ;; Drop current frame, reuse stack space — true TCO
              (dict-set! vm "frames" rest-frames)
              (dict-set! vm "sp" (get frame "base"))
              (vm-call vm f args-rev)))

        (= op 50)  ;; OP_RETURN
          (let ((result (vm-pop vm)))
            (dict-set! vm "frames" rest-frames)
            (dict-set! vm "sp" (get frame "base"))
            (vm-push vm result))

        (= op 51)  ;; OP_CLOSURE
          (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)  ;; OP_CALL_PRIM
          (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)))

        ;; ---- Collections ----
        (= op 64)  ;; OP_LIST
          (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)  ;; OP_DICT
          (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))

        ;; ---- String ops ----
        (= op 144) ;; OP_STR_CONCAT
          (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)))

        ;; ---- Define ----
        (= op 128) ;; OP_DEFINE
          (let ((idx (frame-read-u16 frame))
                (name (nth consts idx)))
            (dict-set! (get vm "globals") name (vm-peek vm)))

        ;; ---- Inline primitives ----
        (= op 160) ;; OP_ADD
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (+ a b)))
        (= op 161) ;; OP_SUB
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (- a b)))
        (= op 162) ;; OP_MUL
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (* a b)))
        (= op 163) ;; OP_DIV
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (/ a b)))
        (= op 164) ;; OP_EQ
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (= a b)))
        (= op 165) ;; OP_LT
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (< a b)))
        (= op 166) ;; OP_GT
          (let ((b (vm-pop vm)) (a (vm-pop vm)))
            (vm-push vm (> a b)))
        (= op 167) ;; OP_NOT
          (vm-push vm (not (vm-pop vm)))
        (= op 168) ;; OP_LEN
          (vm-push vm (len (vm-pop vm)))
        (= op 169) ;; OP_FIRST
          (vm-push vm (first (vm-pop vm)))
        (= op 170) ;; OP_REST
          (vm-push vm (rest (vm-pop vm)))
        (= op 171) ;; OP_NTH
          (let ((n (vm-pop vm)) (coll (vm-pop vm)))
            (vm-push vm (nth coll n)))
        (= op 172) ;; OP_CONS
          (let ((coll (vm-pop vm)) (x (vm-pop vm)))
            (vm-push vm (cons x coll)))
        (= op 173) ;; OP_NEG
          (vm-push vm (- 0 (vm-pop vm)))
        (= op 174) ;; OP_INC
          (vm-push vm (inc (vm-pop vm)))
        (= op 175) ;; OP_DEC
          (vm-push vm (dec (vm-pop vm)))

        :else
          (error (str "VM: unknown opcode " op))))))


;; --------------------------------------------------------------------------
;; 12. Entry points
;; --------------------------------------------------------------------------

;; Execute a closure with arguments — creates a fresh VM.
(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))))

;; Execute a compiled module (top-level bytecode).
(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)))
        ;; Pad local slots
        (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)))))


;; --------------------------------------------------------------------------
;; 13. Platform interface
;; --------------------------------------------------------------------------
;;
;; Each target must provide:
;;
;;   make-vm-stack(size)        → opaque stack (array-like)
;;   vm-stack-get(stack, idx)   → value at index
;;   vm-stack-set!(stack, idx, value) → mutate index
;;   vm-stack-length(stack)     → current capacity
;;   vm-stack-copy!(src, dst, count) → copy first count elements
;;
;;   cek-call(f, args)          → evaluate via CEK machine (fallback)
;;   get-primitive(name)        → look up primitive by name (returns callable)
;;   call-primitive(name, args) → call primitive directly with args list
;;
;;   env-parent(env)            → parent environment or nil
;;   env-has?(env, name)        → boolean
;;   env-get(env, name)         → value
;;   env-set!(env, name, value) → mutate binding