;; lib/prolog/runtime.sx — unification, trail, and (later) solver
;;
;; Phase 2 focus: runtime terms + unification + trail.
;;
;; Term representation at runtime:
;;   atom       ("atom" name)           -- immutable tagged list, same as parse AST
;;   num        ("num" n)                -- likewise
;;   str        ("str" s)                -- likewise
;;   compound   ("compound" fun args)    -- args is a regular list of terms
;;   var        {:tag "var"              -- MUTABLE dict; :binding = nil (unbound) or term
;;               :name "X"
;;               :id <int>
;;               :binding <term-or-nil>}
;;
;; Parse-time ("var" name) tokens must be instantiated into runtime vars
;; before unification. Fresh renaming happens per clause resolution so
;; that two separate calls to the same clause don't share variables.
;;
;; Trail:
;;   {:entries (<var> <var> ...) :len N}  -- stack of vars that got bound
;;   Mark = integer length of the entries list at checkpoint time.
;;   trail-undo-to! pops entries pushed since the mark, nil'ing :binding.
;;
;; Occurs-check: off by default; configurable via trail :occurs-check flag.

;; ── Var id counter ─────────────────────────────────────────────────
(define pl-var-counter {:n 0})

(define
  pl-fresh-id
  (fn
    ()
    (let
      ((n (dict-get pl-var-counter :n)))
      (dict-set! pl-var-counter :n (+ n 1))
      n)))

;; ── Term constructors / predicates ─────────────────────────────────
(define pl-mk-rt-var (fn (name) {:tag "var" :id (pl-fresh-id) :name name :binding nil}))

(define pl-var? (fn (t) (and (dict? t) (= (dict-get t :tag) "var"))))

(define
  pl-atom?
  (fn (t) (and (list? t) (not (empty? t)) (= (first t) "atom"))))

(define
  pl-num?
  (fn (t) (and (list? t) (not (empty? t)) (= (first t) "num"))))

(define
  pl-str?
  (fn (t) (and (list? t) (not (empty? t)) (= (first t) "str"))))

(define
  pl-compound?
  (fn (t) (and (list? t) (not (empty? t)) (= (first t) "compound"))))

(define pl-var-name (fn (v) (dict-get v :name)))
(define pl-var-id (fn (v) (dict-get v :id)))
(define pl-var-binding (fn (v) (dict-get v :binding)))
(define pl-var-bound? (fn (v) (not (nil? (dict-get v :binding)))))

(define pl-atom-name (fn (t) (nth t 1)))
(define pl-num-val (fn (t) (nth t 1)))
(define pl-str-val (fn (t) (nth t 1)))
(define pl-fun (fn (t) (nth t 1)))
(define pl-args (fn (t) (nth t 2)))

;; ── Instantiate parse AST into runtime terms ──────────────────────
;; Walk a parser AST term, replacing ("var" name) occurrences with fresh
;; runtime vars. A name->var dict is threaded so that repeated uses of
;; the same variable within a clause share the same runtime var.
;; "_" is anonymous: each occurrence gets a NEW fresh var (never shared).

(define
  pl-instantiate
  (fn
    (ast var-env)
    (cond
      ((pl-var? ast) ast)
      ((not (list? ast)) ast)
      ((empty? ast) ast)
      ((= (first ast) "var")
        (let
          ((name (nth ast 1)))
          (if
            (= name "_")
            (pl-mk-rt-var "_")
            (if
              (dict-has? var-env name)
              (dict-get var-env name)
              (let ((v (pl-mk-rt-var name))) (dict-set! var-env name v) v)))))
      ((= (first ast) "compound")
        (let
          ((fun (nth ast 1)) (args (nth ast 2)))
          (list
            "compound"
            fun
            (map (fn (a) (pl-instantiate a var-env)) args))))
      ((= (first ast) "clause")
        (let
          ((h (pl-instantiate (nth ast 1) var-env))
            (b (pl-instantiate (nth ast 2) var-env)))
          (list "clause" h b)))
      (true ast))))

(define pl-instantiate-fresh (fn (ast) (pl-instantiate ast {})))

;; ── Walk: follow binding chain ─────────────────────────────────────
(define
  pl-walk
  (fn
    (t)
    (if
      (pl-var? t)
      (if (pl-var-bound? t) (pl-walk (pl-var-binding t)) t)
      t)))

;; Deep-walk: recursively resolve variables inside compound terms.
(define
  pl-walk-deep
  (fn
    (t)
    (let
      ((w (pl-walk t)))
      (if
        (pl-compound? w)
        (list "compound" (pl-fun w) (map pl-walk-deep (pl-args w)))
        w))))

;; ── Trail ─────────────────────────────────────────────────────────
(define pl-mk-trail (fn () {:entries () :len 0 :occurs-check false}))

(define pl-trail-mark (fn (trail) (dict-get trail :len)))

(define
  pl-trail-push!
  (fn
    (trail v)
    (dict-set! trail :entries (cons v (dict-get trail :entries)))
    (dict-set! trail :len (+ 1 (dict-get trail :len)))))

(define
  pl-trail-undo-to!
  (fn
    (trail mark)
    (let
      loop
      ()
      (when
        (> (dict-get trail :len) mark)
        (let
          ((entries (dict-get trail :entries)))
          (let
            ((top (first entries)) (rest (rest entries)))
            (dict-set! top :binding nil)
            (dict-set! trail :entries rest)
            (dict-set! trail :len (- (dict-get trail :len) 1))
            (loop)))))))

;; Bind variable v to term t, recording on the trail.
(define
  pl-bind!
  (fn (v t trail) (dict-set! v :binding t) (pl-trail-push! trail v)))

;; ── Occurs check ──────────────────────────────────────────────────
(define
  pl-occurs?
  (fn
    (v t)
    (let
      ((w (pl-walk t)))
      (cond
        ((pl-var? w) (= (pl-var-id v) (pl-var-id w)))
        ((pl-compound? w) (some (fn (a) (pl-occurs? v a)) (pl-args w)))
        (true false)))))

;; ── Unify ─────────────────────────────────────────────────────────
;; Unify two terms, mutating trail. Returns true on success.
;; On failure, the caller must undo to a pre-unify mark.
(define
  pl-unify!
  (fn
    (t1 t2 trail)
    (let
      ((a (pl-walk t1)) (b (pl-walk t2)))
      (cond
        ((and (pl-var? a) (pl-var? b) (= (pl-var-id a) (pl-var-id b)))
          true)
        ((pl-var? a)
          (if
            (and (dict-get trail :occurs-check) (pl-occurs? a b))
            false
            (do (pl-bind! a b trail) true)))
        ((pl-var? b)
          (if
            (and (dict-get trail :occurs-check) (pl-occurs? b a))
            false
            (do (pl-bind! b a trail) true)))
        ((and (pl-atom? a) (pl-atom? b))
          (= (pl-atom-name a) (pl-atom-name b)))
        ((and (pl-num? a) (pl-num? b)) (= (pl-num-val a) (pl-num-val b)))
        ((and (pl-str? a) (pl-str? b)) (= (pl-str-val a) (pl-str-val b)))
        ((and (pl-compound? a) (pl-compound? b))
          (if
            (and
              (= (pl-fun a) (pl-fun b))
              (= (len (pl-args a)) (len (pl-args b))))
            (pl-unify-lists! (pl-args a) (pl-args b) trail)
            false))
        (true false)))))

(define
  pl-unify-lists!
  (fn
    (xs ys trail)
    (cond
      ((and (empty? xs) (empty? ys)) true)
      ((or (empty? xs) (empty? ys)) false)
      (true
        (if
          (pl-unify! (first xs) (first ys) trail)
          (pl-unify-lists! (rest xs) (rest ys) trail)
          false)))))

;; Convenience: try-unify with auto-undo on failure.
(define
  pl-try-unify!
  (fn
    (t1 t2 trail)
    (let
      ((mark (pl-trail-mark trail)))
      (if
        (pl-unify! t1 t2 trail)
        true
        (do (pl-trail-undo-to! trail mark) false)))))

(define pl-mk-db (fn () {:clauses {}}))

(define
  pl-head-key
  (fn
    (head)
    (cond
      ((pl-compound? head) (str (pl-fun head) "/" (len (pl-args head))))
      ((pl-atom? head) (str (pl-atom-name head) "/0"))
      (true (error "pl-head-key: invalid head")))))

(define pl-clause-key (fn (clause) (pl-head-key (nth clause 1))))

(define pl-goal-key (fn (goal) (pl-head-key goal)))

(define
  pl-db-add!
  (fn
    (db clause)
    (let
      ((key (pl-clause-key clause)) (table (dict-get db :clauses)))
      (cond
        ((nil? (dict-get table key)) (dict-set! table key (list clause)))
        (true (begin (append! (dict-get table key) clause) nil))))))

(define
  pl-db-load!
  (fn
    (db program)
    (cond
      ((empty? program) nil)
      (true
        (begin
          (pl-db-add! db (first program))
          (pl-db-load! db (rest program)))))))

(define
  pl-db-lookup
  (fn
    (db key)
    (let
      ((v (dict-get (dict-get db :clauses) key)))
      (cond ((nil? v) (list)) (true v)))))

(define
  pl-db-lookup-goal
  (fn (db goal) (pl-db-lookup db (pl-goal-key goal))))

(define
  pl-cut?
  (fn (t) (and (list? t) (not (empty? t)) (= (first t) "cut"))))

(define
  pl-solve!
  (fn
    (db goal trail cut-box k)
    (let
      ((g (pl-walk goal)))
      (cond
        ((pl-var? g) false)
        ((pl-cut? g) (begin (dict-set! cut-box :cut true) (k)))
        ((and (pl-atom? g) (= (pl-atom-name g) "true")) (k))
        ((and (pl-atom? g) (= (pl-atom-name g) "fail")) false)
        ((and (pl-atom? g) (= (pl-atom-name g) "nl"))
          (begin (pl-output-write! "\n") (k)))
        ((and (pl-compound? g) (= (pl-fun g) "=") (= (len (pl-args g)) 2))
          (pl-solve-eq! (first (pl-args g)) (nth (pl-args g) 1) trail k))
        ((and (pl-compound? g) (= (pl-fun g) "\\=") (= (len (pl-args g)) 2))
          (pl-solve-not-eq!
            (first (pl-args g))
            (nth (pl-args g) 1)
            trail
            k))
        ((and (pl-compound? g) (= (pl-fun g) "is") (= (len (pl-args g)) 2))
          (pl-solve-eq!
            (first (pl-args g))
            (list "num" (pl-eval-arith (nth (pl-args g) 1)))
            trail
            k))
        ((and (pl-compound? g) (= (pl-fun g) ",") (= (len (pl-args g)) 2))
          (pl-solve!
            db
            (first (pl-args g))
            trail
            cut-box
            (fn () (pl-solve! db (nth (pl-args g) 1) trail cut-box k))))
        ((and (pl-compound? g) (= (pl-fun g) ";") (= (len (pl-args g)) 2))
          (pl-solve-or!
            db
            (first (pl-args g))
            (nth (pl-args g) 1)
            trail
            cut-box
            k))
        ((and (pl-compound? g) (= (pl-fun g) "->") (= (len (pl-args g)) 2))
          (pl-solve-if-then-else!
            db
            (first (pl-args g))
            (nth (pl-args g) 1)
            (list "atom" "fail")
            trail
            cut-box
            k))
        ((and (pl-compound? g) (= (pl-fun g) "call") (= (len (pl-args g)) 1))
          (let
            ((call-cb {:cut false}))
            (pl-solve! db (first (pl-args g)) trail call-cb k)))
        ((and (pl-compound? g) (= (pl-fun g) "write") (= (len (pl-args g)) 1))
          (begin
            (pl-output-write! (pl-format-term (first (pl-args g))))
            (k)))
        (true (pl-solve-user! db g trail cut-box k))))))

(define
  pl-solve-or!
  (fn
    (db a b trail cut-box k)
    (cond
      ((and (pl-compound? a) (= (pl-fun a) "->") (= (len (pl-args a)) 2))
        (pl-solve-if-then-else!
          db
          (first (pl-args a))
          (nth (pl-args a) 1)
          b
          trail
          cut-box
          k))
      (true
        (let
          ((mark (pl-trail-mark trail)))
          (let
            ((r (pl-solve! db a trail cut-box k)))
            (cond
              (r true)
              ((dict-get cut-box :cut) false)
              (true
                (begin
                  (pl-trail-undo-to! trail mark)
                  (pl-solve! db b trail cut-box k))))))))))

(define
  pl-solve-if-then-else!
  (fn
    (db cond-goal then-goal else-goal trail cut-box k)
    (let
      ((mark (pl-trail-mark trail)))
      (let
        ((local-cb {:cut false}))
        (let
          ((found {:val false}))
          (pl-solve!
            db
            cond-goal
            trail
            local-cb
            (fn () (begin (dict-set! found :val true) true)))
          (cond
            ((dict-get found :val) (pl-solve! db then-goal trail cut-box k))
            (true
              (begin
                (pl-trail-undo-to! trail mark)
                (pl-solve! db else-goal trail cut-box k)))))))))

(define pl-output-buffer "")

(define pl-output-clear! (fn () (set! pl-output-buffer "")))

(define
  pl-output-write!
  (fn (s) (begin (set! pl-output-buffer (str pl-output-buffer s)) nil)))

(define
  pl-format-args
  (fn
    (args)
    (cond
      ((empty? args) "")
      ((= (len args) 1) (pl-format-term (first args)))
      (true
        (str
          (pl-format-term (first args))
          ", "
          (pl-format-args (rest args)))))))

(define
  pl-format-term
  (fn
    (t)
    (let
      ((w (pl-walk-deep t)))
      (cond
        ((pl-var? w) (str "_" (pl-var-id w)))
        ((pl-atom? w) (pl-atom-name w))
        ((pl-num? w) (str (pl-num-val w)))
        ((pl-str? w) (pl-str-val w))
        ((pl-compound? w)
          (str (pl-fun w) "(" (pl-format-args (pl-args w)) ")"))
        (true (str w))))))

(define
  pl-eval-arith
  (fn
    (t)
    (let
      ((w (pl-walk-deep t)))
      (cond
        ((pl-num? w) (pl-num-val w))
        ((pl-compound? w)
          (let
            ((f (pl-fun w)) (args (pl-args w)))
            (cond
              ((and (= f "+") (= (len args) 2))
                (+
                  (pl-eval-arith (first args))
                  (pl-eval-arith (nth args 1))))
              ((and (= f "-") (= (len args) 2))
                (-
                  (pl-eval-arith (first args))
                  (pl-eval-arith (nth args 1))))
              ((and (= f "-") (= (len args) 1))
                (- 0 (pl-eval-arith (first args))))
              ((and (= f "*") (= (len args) 2))
                (*
                  (pl-eval-arith (first args))
                  (pl-eval-arith (nth args 1))))
              ((and (= f "/") (= (len args) 2))
                (/
                  (pl-eval-arith (first args))
                  (pl-eval-arith (nth args 1))))
              ((and (= f "mod") (= (len args) 2))
                (mod
                  (pl-eval-arith (first args))
                  (pl-eval-arith (nth args 1))))
              ((and (= f "abs") (= (len args) 1))
                (let
                  ((v (pl-eval-arith (first args))))
                  (cond ((< v 0) (- 0 v)) (true v))))
              (true 0))))
        (true 0)))))

(define
  pl-solve-not-eq!
  (fn
    (a b trail k)
    (let
      ((mark (pl-trail-mark trail)))
      (let
        ((unified (pl-unify! a b trail)))
        (begin
          (pl-trail-undo-to! trail mark)
          (cond (unified false) (true (k))))))))

(define
  pl-solve-eq!
  (fn
    (a b trail k)
    (let
      ((mark (pl-trail-mark trail)))
      (cond
        ((pl-unify! a b trail)
          (let
            ((r (k)))
            (cond
              (r true)
              (true (begin (pl-trail-undo-to! trail mark) false)))))
        (true (begin (pl-trail-undo-to! trail mark) false))))))

(define
  pl-solve-user!
  (fn
    (db goal trail outer-cut-box k)
    (let
      ((inner-cut-box {:cut false}))
      (let
        ((outer-was-cut (dict-get outer-cut-box :cut)))
        (pl-try-clauses!
          db
          goal
          trail
          (pl-db-lookup-goal db goal)
          outer-cut-box
          outer-was-cut
          inner-cut-box
          k)))))

(define
  pl-try-clauses!
  (fn
    (db goal trail clauses outer-cut-box outer-was-cut inner-cut-box k)
    (cond
      ((empty? clauses) false)
      (true
        (let
          ((mark (pl-trail-mark trail)))
          (let
            ((clause (pl-instantiate-fresh (first clauses))))
            (let
              ((head (nth clause 1)) (body (nth clause 2)))
              (cond
                ((pl-unify! goal head trail)
                  (let
                    ((r (pl-solve! db body trail inner-cut-box k)))
                    (cond
                      (r true)
                      ((dict-get inner-cut-box :cut)
                        (begin (pl-trail-undo-to! trail mark) false))
                      ((and (not outer-was-cut) (dict-get outer-cut-box :cut))
                        (begin (pl-trail-undo-to! trail mark) false))
                      (true
                        (begin
                          (pl-trail-undo-to! trail mark)
                          (pl-try-clauses!
                            db
                            goal
                            trail
                            (rest clauses)
                            outer-cut-box
                            outer-was-cut
                            inner-cut-box
                            k))))))
                (true
                  (begin
                    (pl-trail-undo-to! trail mark)
                    (pl-try-clauses!
                      db
                      goal
                      trail
                      (rest clauses)
                      outer-cut-box
                      outer-was-cut
                      inner-cut-box
                      k)))))))))))

(define
  pl-solve-once!
  (fn (db goal trail) (pl-solve! db goal trail {:cut false} (fn () true))))

(define
  pl-solve-count!
  (fn
    (db goal trail)
    (let
      ((box {:n 0}))
      (pl-solve!
        db
        goal
        trail
        {:cut false}
        (fn () (begin (dict-set! box :n (+ (dict-get box :n) 1)) false)))
      (dict-get box :n))))