;; ==========================================================================
;; compiler.sx — SX bytecode compiler
;;
;; Compiles SX AST to bytecode for the platform-native VM.
;; Written in SX — runs on any platform with an SX evaluator.
;;
;; Architecture:
;;   Pass 1: Scope analysis — resolve variables, detect tail positions
;;   Pass 2: Code generation — emit bytecode
;;
;; The compiler produces Code objects (bytecode + constant pool).
;; The VM executes them with a stack machine model.
;; ==========================================================================


;; --------------------------------------------------------------------------
;; Constant pool builder
;; --------------------------------------------------------------------------

(define make-pool
  (fn ()
    {:entries (if (primitive? "mutable-list") (mutable-list) (list))
     :index   {:_count 0}}))

(define pool-add
  (fn (pool value)
    "Add a value to the constant pool, return its index. Deduplicates."
    (let ((key (serialize value))
          (idx-map (get pool "index")))
      (if (has-key? idx-map key)
        (get idx-map key)
        (let ((idx (get idx-map "_count")))
          (dict-set! idx-map key idx)
          (dict-set! idx-map "_count" (+ idx 1))
          (append! (get pool "entries") value)
          idx)))))


;; --------------------------------------------------------------------------
;; Scope analysis
;; --------------------------------------------------------------------------

(define make-scope
  (fn (parent)
    {:locals     (list)   ;; list of {name, slot, mutable?}
     :upvalues   (list)   ;; list of {name, is-local, index}
     :parent     parent
     :next-slot  0}))

(define scope-define-local
  (fn (scope name)
    "Add a local variable, return its slot index."
    (let ((slot (get scope "next-slot")))
      (append! (get scope "locals")
        {:name name :slot slot :mutable false})
      (dict-set! scope "next-slot" (+ slot 1))
      slot)))

(define scope-resolve
  (fn (scope name)
    "Resolve a variable name. Returns {:type \"local\"|\"upvalue\"|\"global\", :index N}."
    (if (nil? scope)
      {:type "global" :index name}
      ;; Check locals
      (let ((locals (get scope "locals"))
            (found (some (fn (l) (= (get l "name") name)) locals)))
        (if found
          (let ((local (first (filter (fn (l) (= (get l "name") name)) locals))))
            {:type "local" :index (get local "slot")})
          ;; Check upvalues (already captured)
          (let ((upvals (get scope "upvalues"))
                (uv-found (some (fn (u) (= (get u "name") name)) upvals)))
            (if uv-found
              (let ((uv (first (filter (fn (u) (= (get u "name") name)) upvals))))
                {:type "upvalue" :index (get uv "index")})
              ;; Try parent scope — if found, capture as upvalue
              (let ((parent-result (scope-resolve (get scope "parent") name)))
                (if (= (get parent-result "type") "global")
                  parent-result
                  ;; Capture from parent as upvalue
                  (let ((uv-idx (len (get scope "upvalues"))))
                    (append! (get scope "upvalues")
                      {:name name
                       :is-local (= (get parent-result "type") "local")
                       :index (get parent-result "index")})
                    {:type "upvalue" :index uv-idx}))))))))))


;; --------------------------------------------------------------------------
;; Code emitter
;; --------------------------------------------------------------------------

(define make-emitter
  (fn ()
    {:bytecode (if (primitive? "mutable-list") (mutable-list) (list))
     :pool     (make-pool)}))

(define emit-byte
  (fn (em byte)
    (append! (get em "bytecode") byte)))

(define emit-u16
  (fn (em value)
    (emit-byte em (mod value 256))
    (emit-byte em (mod (floor (/ value 256)) 256))))

(define emit-i16
  (fn (em value)
    (let ((v (if (< value 0) (+ value 65536) value)))
      (emit-u16 em v))))

(define emit-op
  (fn (em opcode)
    (emit-byte em opcode)))

(define emit-const
  (fn (em value)
    (let ((idx (pool-add (get em "pool") value)))
      (emit-op em 1)   ;; OP_CONST
      (emit-u16 em idx))))

(define current-offset
  (fn (em)
    (len (get em "bytecode"))))

(define patch-i16
  (fn (em offset value)
    "Patch a previously emitted i16 at the given bytecode offset."
    (let ((v (if (< value 0) (+ value 65536) value))
          (bc (get em "bytecode")))
      ;; Direct mutation of bytecode list at offset
      (set-nth! bc offset (mod v 256))
      (set-nth! bc (+ offset 1) (mod (floor (/ v 256)) 256)))))


;; --------------------------------------------------------------------------
;; Compilation — expression dispatch
;; --------------------------------------------------------------------------

(define compile-expr
  (fn (em expr scope tail?)
    "Compile an expression. tail? indicates tail position for TCO."
    (cond
      ;; Nil
      (nil? expr)
        (emit-op em 2)  ;; OP_NIL

      ;; Number
      (= (type-of expr) "number")
        (emit-const em expr)

      ;; String
      (= (type-of expr) "string")
        (emit-const em expr)

      ;; Boolean
      (= (type-of expr) "boolean")
        (emit-op em (if expr 3 4))  ;; OP_TRUE / OP_FALSE

      ;; Keyword
      (= (type-of expr) "keyword")
        (emit-const em (keyword-name expr))

      ;; Symbol — resolve to local/upvalue/global
      (= (type-of expr) "symbol")
        (compile-symbol em (symbol-name expr) scope)

      ;; List — dispatch on head
      (= (type-of expr) "list")
        (if (empty? expr)
          (do (emit-op em 64) (emit-u16 em 0))  ;; OP_LIST 0
          (compile-list em expr scope tail?))

      ;; Dict literal
      (= (type-of expr) "dict")
        (compile-dict em expr scope)

      ;; Fallback
      :else
        (emit-const em expr))))


(define compile-symbol
  (fn (em name scope)
    (let ((resolved (scope-resolve scope name)))
      (cond
        (= (get resolved "type") "local")
          (do (emit-op em 16)  ;; OP_LOCAL_GET
              (emit-byte em (get resolved "index")))
        (= (get resolved "type") "upvalue")
          (do (emit-op em 18)  ;; OP_UPVALUE_GET
              (emit-byte em (get resolved "index")))
        :else
          ;; Global or primitive
          (let ((idx (pool-add (get em "pool") name)))
            (emit-op em 20)  ;; OP_GLOBAL_GET
            (emit-u16 em idx))))))


(define compile-dict
  (fn (em expr scope)
    (let ((ks (keys expr))
          (count (len ks)))
      (for-each (fn (k)
        (emit-const em k)
        (compile-expr em (get expr k) scope false))
        ks)
      (emit-op em 65)     ;; OP_DICT
      (emit-u16 em count))))


;; --------------------------------------------------------------------------
;; List compilation — special forms, calls
;; --------------------------------------------------------------------------

(define compile-list
  (fn (em expr scope tail?)
    (let ((head (first expr))
          (args (rest expr)))
      (if (not (= (type-of head) "symbol"))
        ;; Non-symbol head — compile as call
        (compile-call em head args scope tail?)
        ;; Symbol head — check for special forms
        (let ((name (symbol-name head)))
          (cond
            (= name "if")      (compile-if em args scope tail?)
            (= name "when")    (compile-when em args scope tail?)
            (= name "and")     (compile-and em args scope tail?)
            (= name "or")      (compile-or em args scope tail?)
            (= name "let")     (compile-let em args scope tail?)
            (= name "let*")    (compile-let em args scope tail?)
            (= name "begin")   (compile-begin em args scope tail?)
            (= name "do")      (compile-begin em args scope tail?)
            (= name "lambda")  (compile-lambda em args scope)
            (= name "fn")      (compile-lambda em args scope)
            (= name "define")  (compile-define em args scope)
            (= name "set!")    (compile-set em args scope)
            (= name "quote")   (compile-quote em args)
            (= name "cond")    (compile-cond em args scope tail?)
            (= name "case")    (compile-case em args scope tail?)
            (= name "->")      (compile-thread em args scope tail?)
            (= name "defcomp") (compile-defcomp em args scope)
            (= name "defisland") (compile-defcomp em args scope)
            (= name "defmacro") (compile-defmacro em args scope)
            (= name "defstyle") (emit-op em 2)  ;; defstyle → nil (no-op at runtime)
            (= name "defhandler") (emit-op em 2)  ;; no-op
            (= name "defpage") (emit-op em 2)  ;; handled by page loader
            (= name "defquery") (emit-op em 2)
            (= name "defaction") (emit-op em 2)
            (= name "defrelation") (emit-op em 2)
            (= name "deftype") (emit-op em 2)
            (= name "defeffect") (emit-op em 2)
            (= name "defisland") (compile-defcomp em args scope)
            (= name "quasiquote") (compile-quasiquote em (first args) scope)
            (= name "letrec") (compile-let em args scope tail?)
            ;; Default — function call
            :else
              (compile-call em head args scope tail?)))))))


;; --------------------------------------------------------------------------
;; Special form compilation
;; --------------------------------------------------------------------------

(define compile-if
  (fn (em args scope tail?)
    (let ((test (first args))
          (then-expr (nth args 1))
          (else-expr (if (> (len args) 2) (nth args 2) nil)))
      ;; Compile test
      (compile-expr em test scope false)
      ;; Jump if false to else
      (emit-op em 33)  ;; OP_JUMP_IF_FALSE
      (let ((else-jump (current-offset em)))
        (emit-i16 em 0)  ;; placeholder
        ;; Compile then (in tail position if if is)
        (compile-expr em then-expr scope tail?)
        ;; Jump over else
        (emit-op em 32)  ;; OP_JUMP
        (let ((end-jump (current-offset em)))
          (emit-i16 em 0)  ;; placeholder
          ;; Patch else jump
          (patch-i16 em else-jump (- (current-offset em) (+ else-jump 2)))
          ;; Compile else
          (if (nil? else-expr)
            (emit-op em 2)  ;; OP_NIL
            (compile-expr em else-expr scope tail?))
          ;; Patch end jump
          (patch-i16 em end-jump (- (current-offset em) (+ end-jump 2))))))))


(define compile-when
  (fn (em args scope tail?)
    (let ((test (first args))
          (body (rest args)))
      (compile-expr em test scope false)
      (emit-op em 33)  ;; OP_JUMP_IF_FALSE
      (let ((skip-jump (current-offset em)))
        (emit-i16 em 0)
        (compile-begin em body scope tail?)
        (emit-op em 32)  ;; OP_JUMP
        (let ((end-jump (current-offset em)))
          (emit-i16 em 0)
          (patch-i16 em skip-jump (- (current-offset em) (+ skip-jump 2)))
          (emit-op em 2)  ;; OP_NIL
          (patch-i16 em end-jump (- (current-offset em) (+ end-jump 2))))))))


(define compile-and
  (fn (em args scope tail?)
    (if (empty? args)
      (emit-op em 3)  ;; OP_TRUE
      (if (= (len args) 1)
        (compile-expr em (first args) scope tail?)
        (do
          (compile-expr em (first args) scope false)
          (emit-op em 6)  ;; OP_DUP
          (emit-op em 33)  ;; OP_JUMP_IF_FALSE
          (let ((skip (current-offset em)))
            (emit-i16 em 0)
            (emit-op em 5)  ;; OP_POP (discard duplicated truthy)
            (compile-and em (rest args) scope tail?)
            (patch-i16 em skip (- (current-offset em) (+ skip 2)))))))))


(define compile-or
  (fn (em args scope tail?)
    (if (empty? args)
      (emit-op em 4)  ;; OP_FALSE
      (if (= (len args) 1)
        (compile-expr em (first args) scope tail?)
        (do
          (compile-expr em (first args) scope false)
          (emit-op em 6)  ;; OP_DUP
          (emit-op em 34)  ;; OP_JUMP_IF_TRUE
          (let ((skip (current-offset em)))
            (emit-i16 em 0)
            (emit-op em 5)  ;; OP_POP
            (compile-or em (rest args) scope tail?)
            (patch-i16 em skip (- (current-offset em) (+ skip 2)))))))))


(define compile-begin
  (fn (em exprs scope tail?)
    (if (empty? exprs)
      (emit-op em 2)  ;; OP_NIL
      (if (= (len exprs) 1)
        (compile-expr em (first exprs) scope tail?)
        (do
          (compile-expr em (first exprs) scope false)
          (emit-op em 5)  ;; OP_POP
          (compile-begin em (rest exprs) scope tail?))))))


(define compile-let
  (fn (em args scope tail?)
    (let ((bindings (first args))
          (body (rest args))
          (let-scope (make-scope scope)))
      ;; Compile each binding
      (for-each (fn (binding)
        (let ((name (if (= (type-of (first binding)) "symbol")
                      (symbol-name (first binding))
                      (first binding)))
              (value (nth binding 1))
              (slot (scope-define-local let-scope name)))
          (compile-expr em value let-scope false)
          (emit-op em 17)  ;; OP_LOCAL_SET
          (emit-byte em slot)))
        bindings)
      ;; Compile body in let scope
      (compile-begin em body let-scope tail?))))


(define compile-lambda
  (fn (em args scope)
    (let ((params (first args))
          (body (rest args))
          (fn-scope (make-scope scope))
          (fn-em (make-emitter)))
      ;; Define params as locals in fn scope
      (for-each (fn (p)
        (let ((name (if (= (type-of p) "symbol") (symbol-name p) p)))
          (when (and (not (= name "&key"))
                     (not (= name "&rest")))
            (scope-define-local fn-scope name))))
        params)
      ;; Compile body
      (compile-begin fn-em body fn-scope true)  ;; tail position
      (emit-op fn-em 50)  ;; OP_RETURN
      ;; Add code object to parent constant pool
      (let ((code {:arity (len (get fn-scope "locals"))
                   :bytecode (get fn-em "bytecode")
                   :constants (get (get fn-em "pool") "entries")
                   :upvalues (get fn-scope "upvalues")})
            (code-idx (pool-add (get em "pool") code)))
        (emit-op em 51)  ;; OP_CLOSURE
        (emit-u16 em code-idx)))))


(define compile-define
  (fn (em args scope)
    (let ((name-expr (first args))
          (name (if (= (type-of name-expr) "symbol")
                  (symbol-name name-expr)
                  name-expr))
          (value (nth args 1))
          (name-idx (pool-add (get em "pool") name)))
      (compile-expr em value scope false)
      (emit-op em 128)  ;; OP_DEFINE
      (emit-u16 em name-idx))))


(define compile-set
  (fn (em args scope)
    (let ((name (if (= (type-of (first args)) "symbol")
                  (symbol-name (first args))
                  (first args)))
          (value (nth args 1))
          (resolved (scope-resolve scope name)))
      (compile-expr em value scope false)
      (cond
        (= (get resolved "type") "local")
          (do (emit-op em 17)  ;; OP_LOCAL_SET
              (emit-byte em (get resolved "index")))
        (= (get resolved "type") "upvalue")
          (do (emit-op em 19)  ;; OP_UPVALUE_SET
              (emit-byte em (get resolved "index")))
        :else
          (let ((idx (pool-add (get em "pool") name)))
            (emit-op em 21)  ;; OP_GLOBAL_SET
            (emit-u16 em idx))))))


(define compile-quote
  (fn (em args)
    (if (empty? args)
      (emit-op em 2)  ;; OP_NIL
      (emit-const em (first args)))))


(define compile-cond
  (fn (em args scope tail?)
    "Compile (cond test1 body1 test2 body2 ... :else fallback)."
    (if (< (len args) 2)
      (emit-op em 2)  ;; OP_NIL
      (let ((test (first args))
            (body (nth args 1))
            (rest-clauses (if (> (len args) 2) (slice args 2) (list))))
        (if (or (= test :else) (= test true))
          ;; else clause — just compile the body
          (compile-expr em body scope tail?)
          (do
            (compile-expr em test scope false)
            (emit-op em 33)  ;; OP_JUMP_IF_FALSE
            (let ((skip (current-offset em)))
              (emit-i16 em 0)
              (compile-expr em body scope tail?)
              (emit-op em 32)  ;; OP_JUMP
              (let ((end-jump (current-offset em)))
                (emit-i16 em 0)
                (patch-i16 em skip (- (current-offset em) (+ skip 2)))
                (compile-cond em rest-clauses scope tail?)
                (patch-i16 em end-jump (- (current-offset em) (+ end-jump 2)))))))))))


(define compile-case
  (fn (em args scope tail?)
    "Compile (case expr val1 body1 val2 body2 ... :else fallback)."
    ;; Desugar to nested if: evaluate expr once, then compare
    (compile-expr em (first args) scope false)
    (let ((clauses (rest args)))
      (compile-case-clauses em clauses scope tail?))))

(define compile-case-clauses
  (fn (em clauses scope tail?)
    (if (< (len clauses) 2)
      (do (emit-op em 5) (emit-op em 2))  ;; POP match-val, push NIL
      (let ((test (first clauses))
            (body (nth clauses 1))
            (rest-clauses (if (> (len clauses) 2) (slice clauses 2) (list))))
        (if (or (= test :else) (= test true))
          (do (emit-op em 5)  ;; POP match-val
              (compile-expr em body scope tail?))
          (do
            (emit-op em 6)  ;; DUP match-val
            (compile-expr em test scope false)
            (let ((name-idx (pool-add (get em "pool") "=")))
              (emit-op em 52) (emit-u16 em name-idx) (emit-byte em 2))  ;; CALL_PRIM "=" 2
            (emit-op em 33)  ;; JUMP_IF_FALSE
            (let ((skip (current-offset em)))
              (emit-i16 em 0)
              (emit-op em 5)  ;; POP match-val
              (compile-expr em body scope tail?)
              (emit-op em 32)  ;; JUMP
              (let ((end-jump (current-offset em)))
                (emit-i16 em 0)
                (patch-i16 em skip (- (current-offset em) (+ skip 2)))
                (compile-case-clauses em rest-clauses scope tail?)
                (patch-i16 em end-jump (- (current-offset em) (+ end-jump 2)))))))))))


(define compile-thread
  (fn (em args scope tail?)
    "Compile (-> val (f1 a) (f2 b)) by desugaring to nested calls."
    (if (empty? args)
      (emit-op em 2)
      (if (= (len args) 1)
        (compile-expr em (first args) scope tail?)
        ;; Desugar: (-> x (f a)) → (f x a)
        (let ((val-expr (first args))
              (forms (rest args)))
          (compile-thread-step em val-expr forms scope tail?))))))

(define compile-thread-step
  (fn (em val-expr forms scope tail?)
    (if (empty? forms)
      (compile-expr em val-expr scope tail?)
      (let ((form (first forms))
            (rest-forms (rest forms))
            (is-tail (and tail? (empty? rest-forms))))
        ;; Build desugared call: (f val args...)
        (let ((call-expr
                (if (list? form)
                  ;; (-> x (f a b)) → (f x a b)
                  (concat (list (first form) val-expr) (rest form))
                  ;; (-> x f) → (f x)
                  (list form val-expr))))
          (if (empty? rest-forms)
            (compile-expr em call-expr scope is-tail)
            (do
              (compile-expr em call-expr scope false)
              ;; Thread result through remaining forms
              ;; Store in temp, compile next step
              ;; Actually, just compile sequentially — each step returns a value
              (compile-thread-step em call-expr rest-forms scope tail?))))))))


(define compile-defcomp
  (fn (em args scope)
    "Compile defcomp/defisland — delegates to runtime primitive."
    ;; For now, quote the entire defcomp form and call the primitive
    (emit-const em (concat (list (make-symbol "defcomp")) args))
    (let ((name-idx (pool-add (get em "pool") "eval-defcomp")))
      (emit-op em 52) (emit-u16 em name-idx) (emit-byte em 1))  ;; CALL_PRIM
    ))

(define compile-defmacro
  (fn (em args scope)
    "Compile defmacro — delegates to runtime."
    (emit-const em (concat (list (make-symbol "defmacro")) args))
    (let ((name-idx (pool-add (get em "pool") "eval-defmacro")))
      (emit-op em 52) (emit-u16 em name-idx) (emit-byte em 1))))


(define compile-quasiquote
  (fn (em expr scope)
    "Compile quasiquote — for now, quote the template and call qq-expand at runtime."
    (emit-const em expr)  ;; push template
    ;; qq-expand is a runtime function — call it
    (let ((name-idx (pool-add (get em "pool") "qq-expand-runtime")))
      (emit-op em 52) (emit-u16 em name-idx) (emit-byte em 1))))


;; --------------------------------------------------------------------------
;; Function call compilation
;; --------------------------------------------------------------------------

(define compile-call
  (fn (em head args scope tail?)
    ;; Check for known primitives
    (let ((is-prim (and (= (type-of head) "symbol")
                        (let ((name (symbol-name head)))
                          (and (not (= (get (scope-resolve scope name) "type") "local"))
                               (not (= (get (scope-resolve scope name) "type") "upvalue"))
                               (primitive? name))))))
      (if is-prim
        ;; Direct primitive call — no closure overhead
        (let ((name (symbol-name head))
              (name-idx (pool-add (get em "pool") name)))
          (for-each (fn (a) (compile-expr em a scope false)) args)
          (emit-op em 52)  ;; OP_CALL_PRIM
          (emit-u16 em name-idx)
          (emit-byte em (len args)))
        ;; General call
        (do
          (compile-expr em head scope false)
          (for-each (fn (a) (compile-expr em a scope false)) args)
          (if tail?
            (do (emit-op em 49)  ;; OP_TAIL_CALL
                (emit-byte em (len args)))
            (do (emit-op em 48)  ;; OP_CALL
                (emit-byte em (len args)))))))))


;; --------------------------------------------------------------------------
;; Top-level API
;; --------------------------------------------------------------------------

(define compile
  (fn (expr)
    "Compile a single SX expression to a bytecode module."
    (let ((em (make-emitter))
          (scope (make-scope nil)))
      (compile-expr em expr scope false)
      (emit-op em 50)  ;; OP_RETURN
      {:bytecode (get em "bytecode")
       :constants (get (get em "pool") "entries")})))

(define compile-module
  (fn (exprs)
    "Compile a list of top-level expressions to a bytecode module."
    (let ((em (make-emitter))
          (scope (make-scope nil)))
      (for-each (fn (expr)
        (compile-expr em expr scope false)
        (emit-op em 5))  ;; OP_POP between top-level exprs
        (init exprs))
      ;; Last expression's value is the module result
      (compile-expr em (last exprs) scope false)
      (emit-op em 50)  ;; OP_RETURN
      {:bytecode (get em "bytecode")
       :constants (get (get em "pool") "entries")})))