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