;; Smalltalk AST evaluator — sequential semantics. Method dispatch uses the
;; class table from runtime.sx; native receivers fall back to a primitive
;; method table. Non-local return is implemented via a sentinel marker; the
;; full continuation-based escape is the Phase 3 showcase.
;;
;; Frame:
;;   {:self V                 ; receiver
;;    :method-class N         ; defining class of the executing method
;;    :locals (mutable dict)  ; param + temp bindings
;;    :parent P}              ; outer frame for blocks (nil for top-level)
;;
;; `smalltalk-eval-ast(ast, frame)` returns the value or a return marker.
;; Method invocation unwraps return markers; sequences propagate them.

(define
  st-make-frame
  (fn
    (self method-class parent)
    {:self self :method-class method-class :locals {} :parent parent}))

(define st-return-marker (fn (v) {:st-return true :value v}))

(define
  st-return-marker?
  (fn (v) (and (dict? v) (has-key? v :st-return) (= (get v :st-return) true))))

(define
  st-make-block
  (fn
    (ast frame)
    {:type "st-block"
     :params (get ast :params)
     :temps (get ast :temps)
     :body (get ast :body)
     :env frame}))

(define
  st-block?
  (fn
    (v)
    (and (dict? v) (has-key? v :type) (= (get v :type) "st-block"))))

(define
  st-class-ref
  (fn (name) {:type "st-class" :name name}))

(define
  st-class-ref?
  (fn (v) (and (dict? v) (has-key? v :type) (= (get v :type) "st-class"))))

;; Walk the frame chain looking for a local binding.
(define
  st-lookup-local
  (fn
    (frame name)
    (cond
      ((= frame nil) {:found false :value nil :frame nil})
      ((has-key? (get frame :locals) name)
        {:found true :value (get (get frame :locals) name) :frame frame})
      (else (st-lookup-local (get frame :parent) name)))))

;; Walk the frame chain looking for the frame whose self has this ivar.
(define
  st-lookup-ivar-frame
  (fn
    (frame name)
    (cond
      ((= frame nil) nil)
      ((let ((self (get frame :self)))
         (and (st-instance? self) (has-key? (get self :ivars) name)))
        frame)
      (else (st-lookup-ivar-frame (get frame :parent) name)))))

;; Resolve an identifier in eval order: local → ivar → class → error.
(define
  st-resolve-ident
  (fn
    (frame name)
    (let
      ((local-result (st-lookup-local frame name)))
      (cond
        ((get local-result :found) (get local-result :value))
        (else
          (let
            ((iv-frame (st-lookup-ivar-frame frame name)))
            (cond
              ((not (= iv-frame nil))
                (get (get (get iv-frame :self) :ivars) name))
              ((st-class-exists? name) (st-class-ref name))
              (else
                (error
                  (str "smalltalk-eval-ast: undefined variable '" name "'"))))))))))

;; Assign to an existing local in the frame chain or, failing that, an ivar
;; on self. Errors if neither exists.
(define
  st-assign!
  (fn
    (frame name value)
    (let
      ((local-result (st-lookup-local frame name)))
      (cond
        ((get local-result :found)
          (begin
            (dict-set! (get (get local-result :frame) :locals) name value)
            value))
        (else
          (let
            ((iv-frame (st-lookup-ivar-frame frame name)))
            (cond
              ((not (= iv-frame nil))
                (begin
                  (dict-set! (get (get iv-frame :self) :ivars) name value)
                  value))
              (else
                ;; Smalltalk allows new locals to be introduced; for our subset
                ;; we treat unknown writes as errors so test mistakes surface.
                (error
                  (str "smalltalk-eval-ast: cannot assign undefined '" name "'"))))))))))

;; ── Main evaluator ─────────────────────────────────────────────────────
(define
  smalltalk-eval-ast
  (fn
    (ast frame)
    (cond
      ((not (dict? ast)) (error (str "smalltalk-eval-ast: bad ast " ast)))
      (else
        (let
          ((ty (get ast :type)))
          (cond
            ((= ty "lit-int") (get ast :value))
            ((= ty "lit-float") (get ast :value))
            ((= ty "lit-string") (get ast :value))
            ((= ty "lit-char") (get ast :value))
            ((= ty "lit-symbol") (make-symbol (get ast :value)))
            ((= ty "lit-nil") nil)
            ((= ty "lit-true") true)
            ((= ty "lit-false") false)
            ((= ty "lit-array")
              (map
                (fn (e) (smalltalk-eval-ast e frame))
                (get ast :elements)))
            ((= ty "lit-byte-array") (get ast :elements))
            ((= ty "self") (get frame :self))
            ((= ty "super") (get frame :self))
            ((= ty "thisContext") frame)
            ((= ty "ident") (st-resolve-ident frame (get ast :name)))
            ((= ty "assign")
              (st-assign! frame (get ast :name) (smalltalk-eval-ast (get ast :expr) frame)))
            ((= ty "return")
              (st-return-marker (smalltalk-eval-ast (get ast :expr) frame)))
            ((= ty "block") (st-make-block ast frame))
            ((= ty "seq") (st-eval-seq (get ast :exprs) frame))
            ((= ty "send")
              (st-eval-send ast frame (= (get (get ast :receiver) :type) "super")))
            ((= ty "cascade") (st-eval-cascade ast frame))
            (else (error (str "smalltalk-eval-ast: unknown type '" ty "'")))))))))

(define
  st-eval-seq
  (fn
    (exprs frame)
    (let ((result nil))
      (begin
        (define
          sq-loop
          (fn
            (rest)
            (cond
              ((= (len rest) 0) nil)
              (else
                (let ((v (smalltalk-eval-ast (nth rest 0) frame)))
                  (cond
                    ((st-return-marker? v) (set! result v))
                    ((= (len rest) 1) (set! result v))
                    (else (sq-loop (rest-of rest)))))))))
        (sq-loop exprs)
        result))))

(define
  rest-of
  (fn
    (lst)
    (let ((out (list)) (i 1) (n (len lst)))
      (begin
        (define
          ro-loop
          (fn
            ()
            (when
              (< i n)
              (begin (append! out (nth lst i)) (set! i (+ i 1)) (ro-loop)))))
        (ro-loop)
        out))))

(define
  st-eval-send
  (fn
    (ast frame super?)
    (let
      ((receiver (smalltalk-eval-ast (get ast :receiver) frame))
        (selector (get ast :selector))
        (args (map (fn (a) (smalltalk-eval-ast a frame)) (get ast :args))))
      (cond
        (super?
          (st-super-send (get frame :self) selector args (get frame :method-class)))
        (else (st-send receiver selector args))))))

(define
  st-eval-cascade
  (fn
    (ast frame)
    (let
      ((receiver (smalltalk-eval-ast (get ast :receiver) frame))
        (msgs (get ast :messages))
        (last nil))
      (begin
        (for-each
          (fn
            (m)
            (let
              ((sel (get m :selector))
                (args (map (fn (a) (smalltalk-eval-ast a frame)) (get m :args))))
              (set! last (st-send receiver sel args))))
          msgs)
        last))))

;; ── Send dispatch ──────────────────────────────────────────────────────
(define
  st-send
  (fn
    (receiver selector args)
    (let
      ((cls (st-class-of-for-send receiver)))
      (let
        ((class-side? (st-class-ref? receiver))
          (recv-class (if (st-class-ref? receiver) (get receiver :name) cls)))
        (let
          ((method
             (if class-side?
               (st-method-lookup recv-class selector true)
               (st-method-lookup recv-class selector false))))
          (cond
            ((not (= method nil))
              (st-invoke method receiver args))
            ((st-block? receiver)
              (let ((bd (st-block-dispatch receiver selector args)))
                (cond
                  ((= bd :unhandled) (st-dnu receiver selector args))
                  (else bd))))
            (else
              (let ((primitive-result (st-primitive-send receiver selector args)))
                (cond
                  ((= primitive-result :unhandled)
                    (st-dnu receiver selector args))
                  (else primitive-result))))))))))

;; Construct a Message object for doesNotUnderstand:.
(define
  st-make-message
  (fn
    (selector args)
    (let ((msg (st-make-instance "Message")))
      (begin
        (dict-set! (get msg :ivars) "selector" (make-symbol selector))
        (dict-set! (get msg :ivars) "arguments" args)
        msg))))

;; Trigger doesNotUnderstand:. If the receiver's class chain defines an
;; override, invoke it with a freshly-built Message; otherwise raise.
(define
  st-dnu
  (fn
    (receiver selector args)
    (let
      ((cls (st-class-of-for-send receiver))
        (class-side? (st-class-ref? receiver)))
      (let
        ((recv-class (if class-side? (get receiver :name) cls)))
        (let
          ((method (st-method-lookup recv-class "doesNotUnderstand:" class-side?)))
          (cond
            ((not (= method nil))
              (let ((msg (st-make-message selector args)))
                (st-invoke method receiver (list msg))))
            (else
              (error
                (str "doesNotUnderstand: " recv-class " >> " selector)))))))))

(define
  st-class-of-for-send
  (fn
    (v)
    (cond
      ((st-class-ref? v) "Class")
      (else (st-class-of v)))))

;; super send: lookup starts at the *defining* class's superclass, not the
;; receiver class. This is what makes inherited methods compose correctly
;; under refinement — a method on Foo that calls `super bar` resolves to
;; Foo's superclass's `bar` regardless of the dynamic receiver class.
(define
  st-super-send
  (fn
    (receiver selector args defining-class)
    (cond
      ((= defining-class nil)
        (error (str "super send outside method context: " selector)))
      (else
        (let
          ((super (st-class-superclass defining-class))
            (class-side? (st-class-ref? receiver)))
          (cond
            ((= super nil)
              (error (str "super send past root: " selector " in " defining-class)))
            (else
              (let ((method (st-method-lookup super selector class-side?)))
                (cond
                  ((not (= method nil))
                    (st-invoke method receiver args))
                  (else
                    ;; Try primitives starting from super's perspective too —
                    ;; for native receivers the primitive table is global, so
                    ;; super basically reaches the same primitives. The point
                    ;; of super is to skip user overrides on the receiver's
                    ;; class chain below `super`, which method-lookup above
                    ;; already enforces.
                    (let ((p (st-primitive-send receiver selector args)))
                      (cond
                        ((= p :unhandled)
                          (st-dnu receiver selector args))
                        (else p)))))))))))))

;; ── Method invocation ──────────────────────────────────────────────────
(define
  st-invoke
  (fn
    (method receiver args)
    (let
      ((params (get method :params))
        (temps (get method :temps))
        (body (get method :body))
        (defining-class (get method :defining-class)))
      (cond
        ((not (= (len params) (len args)))
          (error
            (str "smalltalk-eval-ast: arity mismatch for "
              (get method :selector)
              " expected " (len params) " got " (len args))))
        (else
          (let
            ((frame (st-make-frame receiver defining-class nil)))
            (begin
              ;; Bind params
              (let ((i 0))
                (begin
                  (define
                    pb-loop
                    (fn
                      ()
                      (when
                        (< i (len params))
                        (begin
                          (dict-set!
                            (get frame :locals)
                            (nth params i)
                            (nth args i))
                          (set! i (+ i 1))
                          (pb-loop)))))
                  (pb-loop)))
              ;; Bind temps to nil
              (for-each
                (fn (t) (dict-set! (get frame :locals) t nil))
                temps)
              ;; Execute body
              (let ((result (st-eval-seq body frame)))
                (cond
                  ((st-return-marker? result) (get result :value))
                  (else receiver))))))))))

;; ── Block dispatch ─────────────────────────────────────────────────────
(define
  st-block-value-selector?
  (fn
    (s)
    (or
      (= s "value")
      (= s "value:")
      (= s "value:value:")
      (= s "value:value:value:")
      (= s "value:value:value:value:"))))

(define
  st-block-dispatch
  (fn
    (block selector args)
    (cond
      ((st-block-value-selector? selector) (st-block-apply block args))
      ((= selector "valueWithArguments:") (st-block-apply block (nth args 0)))
      ((= selector "whileTrue:")
        (st-block-while block (nth args 0) true))
      ((= selector "whileFalse:")
        (st-block-while block (nth args 0) false))
      ((= selector "whileTrue") (st-block-while block nil true))
      ((= selector "whileFalse") (st-block-while block nil false))
      ((= selector "numArgs") (len (get block :params)))
      ((= selector "class") (st-class-ref "BlockClosure"))
      ((= selector "==") (= block (nth args 0)))
      ((= selector "printString") "a BlockClosure")
      (else :unhandled))))

(define
  st-block-apply
  (fn
    (block args)
    (let
      ((params (get block :params))
        (temps (get block :temps))
        (body (get block :body))
        (env (get block :env)))
      (cond
        ((not (= (len params) (len args)))
          (error
            (str "BlockClosure: arity mismatch — block expects "
              (len params) " got " (len args))))
        (else
          (let
            ((frame (st-make-frame
                     (if (= env nil) nil (get env :self))
                     (if (= env nil) nil (get env :method-class))
                     env)))
            (begin
              (let ((i 0))
                (begin
                  (define
                    pb-loop
                    (fn
                      ()
                      (when
                        (< i (len params))
                        (begin
                          (dict-set!
                            (get frame :locals)
                            (nth params i)
                            (nth args i))
                          (set! i (+ i 1))
                          (pb-loop)))))
                  (pb-loop)))
              (for-each
                (fn (t) (dict-set! (get frame :locals) t nil))
                temps)
              (st-eval-seq body frame))))))))

(define
  st-block-while
  (fn
    (cond-block body-block target)
    (let ((last nil))
      (begin
        (define
          wh-loop
          (fn
            ()
            (let
              ((c (st-block-apply cond-block (list))))
              (when
                (= c target)
                (begin
                  (cond
                    ((not (= body-block nil))
                      (set! last (st-block-apply body-block (list)))))
                  (wh-loop))))))
        (wh-loop)
        last))))

;; ── Primitive method table for native receivers ────────────────────────
;; Returns the result, or the sentinel :unhandled if no primitive matches —
;; in which case st-send falls back to doesNotUnderstand:.
(define
  st-primitive-send
  (fn
    (receiver selector args)
    (let ((cls (st-class-of receiver)))
      (cond
        ((or (= cls "SmallInteger") (= cls "Float"))
          (st-num-send receiver selector args))
        ((or (= cls "String") (= cls "Symbol"))
          (st-string-send receiver selector args))
        ((= cls "True") (st-bool-send true selector args))
        ((= cls "False") (st-bool-send false selector args))
        ((= cls "UndefinedObject") (st-nil-send selector args))
        ((= cls "Array") (st-array-send receiver selector args))
        ((st-class-ref? receiver) (st-class-side-send receiver selector args))
        (else :unhandled)))))

(define
  st-num-send
  (fn
    (n selector args)
    (cond
      ((= selector "+") (+ n (nth args 0)))
      ((= selector "-") (- n (nth args 0)))
      ((= selector "*") (* n (nth args 0)))
      ((= selector "/") (/ n (nth args 0)))
      ((= selector "//") (/ n (nth args 0)))
      ((= selector "\\\\") (mod n (nth args 0)))
      ((= selector "<") (< n (nth args 0)))
      ((= selector ">") (> n (nth args 0)))
      ((= selector "<=") (<= n (nth args 0)))
      ((= selector ">=") (>= n (nth args 0)))
      ((= selector "=") (= n (nth args 0)))
      ((= selector "~=") (not (= n (nth args 0))))
      ((= selector "==") (= n (nth args 0)))
      ((= selector "~~") (not (= n (nth args 0))))
      ((= selector "negated") (- 0 n))
      ((= selector "abs") (if (< n 0) (- 0 n) n))
      ((= selector "max:") (if (> n (nth args 0)) n (nth args 0)))
      ((= selector "min:") (if (< n (nth args 0)) n (nth args 0)))
      ((= selector "printString") (str n))
      ((= selector "asString") (str n))
      ((= selector "class")
        (st-class-ref (st-class-of n)))
      ((= selector "isNil") false)
      ((= selector "notNil") true)
      ((= selector "isZero") (= n 0))
      ((= selector "between:and:")
        (and (>= n (nth args 0)) (<= n (nth args 1))))
      ((= selector "to:do:")
        (let ((i n) (stop (nth args 0)) (block (nth args 1)))
          (begin
            (define
              td-loop
              (fn
                ()
                (when
                  (<= i stop)
                  (begin
                    (st-block-apply block (list i))
                    (set! i (+ i 1))
                    (td-loop)))))
            (td-loop)
            n)))
      ((= selector "timesRepeat:")
        (let ((i 0) (block (nth args 0)))
          (begin
            (define
              tr-loop
              (fn
                ()
                (when
                  (< i n)
                  (begin
                    (st-block-apply block (list))
                    (set! i (+ i 1))
                    (tr-loop)))))
            (tr-loop)
            n)))
      (else :unhandled))))

(define
  st-string-send
  (fn
    (s selector args)
    (cond
      ((= selector ",") (str s (nth args 0)))
      ((= selector "size") (len s))
      ((= selector "=") (= s (nth args 0)))
      ((= selector "~=") (not (= s (nth args 0))))
      ((= selector "==") (= s (nth args 0)))
      ((= selector "~~") (not (= s (nth args 0))))
      ((= selector "isEmpty") (= (len s) 0))
      ((= selector "notEmpty") (> (len s) 0))
      ((= selector "printString") (str "'" s "'"))
      ((= selector "asString") s)
      ((= selector "asSymbol") (make-symbol (if (symbol? s) (str s) s)))
      ((= selector "class") (st-class-ref (st-class-of s)))
      ((= selector "isNil") false)
      ((= selector "notNil") true)
      (else :unhandled))))

(define
  st-bool-send
  (fn
    (b selector args)
    (cond
      ((= selector "not") (not b))
      ((= selector "&") (and b (nth args 0)))
      ((= selector "|") (or b (nth args 0)))
      ((= selector "and:")
        (cond (b (st-block-apply (nth args 0) (list))) (else false)))
      ((= selector "or:")
        (cond (b true) (else (st-block-apply (nth args 0) (list)))))
      ((= selector "ifTrue:")
        (cond (b (st-block-apply (nth args 0) (list))) (else nil)))
      ((= selector "ifFalse:")
        (cond (b nil) (else (st-block-apply (nth args 0) (list)))))
      ((= selector "ifTrue:ifFalse:")
        (cond
          (b (st-block-apply (nth args 0) (list)))
          (else (st-block-apply (nth args 1) (list)))))
      ((= selector "ifFalse:ifTrue:")
        (cond
          (b (st-block-apply (nth args 1) (list)))
          (else (st-block-apply (nth args 0) (list)))))
      ((= selector "=") (= b (nth args 0)))
      ((= selector "~=") (not (= b (nth args 0))))
      ((= selector "==") (= b (nth args 0)))
      ((= selector "printString") (if b "true" "false"))
      ((= selector "class") (st-class-ref (if b "True" "False")))
      ((= selector "isNil") false)
      ((= selector "notNil") true)
      (else :unhandled))))

(define
  st-nil-send
  (fn
    (selector args)
    (cond
      ((= selector "isNil") true)
      ((= selector "notNil") false)
      ((= selector "ifNil:") (st-block-apply (nth args 0) (list)))
      ((= selector "ifNotNil:") nil)
      ((= selector "ifNil:ifNotNil:") (st-block-apply (nth args 0) (list)))
      ((= selector "ifNotNil:ifNil:") (st-block-apply (nth args 1) (list)))
      ((= selector "=") (= nil (nth args 0)))
      ((= selector "~=") (not (= nil (nth args 0))))
      ((= selector "==") (= nil (nth args 0)))
      ((= selector "printString") "nil")
      ((= selector "class") (st-class-ref "UndefinedObject"))
      (else :unhandled))))

(define
  st-array-send
  (fn
    (a selector args)
    (cond
      ((= selector "size") (len a))
      ((= selector "at:")
        ;; 1-indexed
        (nth a (- (nth args 0) 1)))
      ((= selector "at:put:")
        (begin
          (set-nth! a (- (nth args 0) 1) (nth args 1))
          (nth args 1)))
      ((= selector "first") (nth a 0))
      ((= selector "last") (nth a (- (len a) 1)))
      ((= selector "isEmpty") (= (len a) 0))
      ((= selector "notEmpty") (> (len a) 0))
      ((= selector "do:")
        (begin
          (for-each
            (fn (e) (st-block-apply (nth args 0) (list e)))
            a)
          a))
      ((= selector "collect:")
        (map (fn (e) (st-block-apply (nth args 0) (list e))) a))
      ((= selector "select:")
        (filter (fn (e) (st-block-apply (nth args 0) (list e))) a))
      ((= selector ",")
        (let ((out (list)))
          (begin
            (for-each (fn (e) (append! out e)) a)
            (for-each (fn (e) (append! out e)) (nth args 0))
            out)))
      ((= selector "=") (= a (nth args 0)))
      ((= selector "==") (= a (nth args 0)))
      ((= selector "printString")
        (str "#(" (join " " (map (fn (e) (str e)) a)) ")"))
      ((= selector "class") (st-class-ref "Array"))
      ((= selector "isNil") false)
      ((= selector "notNil") true)
      (else :unhandled))))

(define
  st-class-side-send
  (fn
    (cref selector args)
    (let ((name (get cref :name)))
      (cond
        ((= selector "new") (st-make-instance name))
        ((= selector "name") name)
        ((= selector "superclass")
          (let ((s (st-class-superclass name)))
            (cond ((= s nil) nil) (else (st-class-ref s)))))
        ((= selector "printString") name)
        ((= selector "class") (st-class-ref "Metaclass"))
        ((= selector "==") (and (st-class-ref? (nth args 0))
                                (= name (get (nth args 0) :name))))
        ((= selector "=") (and (st-class-ref? (nth args 0))
                                (= name (get (nth args 0) :name))))
        ((= selector "isNil") false)
        ((= selector "notNil") true)
        (else :unhandled)))))

;; Convenience: parse and evaluate a Smalltalk expression with no receiver.
(define
  smalltalk-eval
  (fn
    (src)
    (let
      ((ast (st-parse-expr src))
        (frame (st-make-frame nil nil nil)))
      (smalltalk-eval-ast ast frame))))

;; Evaluate a sequence of statements at the top level.
(define
  smalltalk-eval-program
  (fn
    (src)
    (let
      ((ast (st-parse src)) (frame (st-make-frame nil nil nil)))
      (begin
        (when
          (and (dict? ast) (has-key? ast :temps))
          (for-each
            (fn (t) (dict-set! (get frame :locals) t nil))
            (get ast :temps)))
        (let ((result (smalltalk-eval-ast ast frame)))
          (cond
            ((st-return-marker? result) (get result :value))
            (else result)))))))