;; lib/prolog/tests/unify.sx — unification + trail unit tests
;;
;; Run via MCP: (pl-unify-tests-run!)
;;
;; Covers: atoms, vars, numbers, strings, compounds, nested compounds,
;; cons-cell lists, trail undo, occurs-check, deep walks.

(define pl-u-count 0)
(define pl-u-pass 0)
(define pl-u-fail 0)
(define pl-u-failures (list))

(define
  pl-u-test!
  (fn
    (name thunk expected)
    (set! pl-u-count (+ pl-u-count 1))
    (let
      ((got (thunk)))
      (if
        (= got expected)
        (set! pl-u-pass (+ pl-u-pass 1))
        (do
          (set! pl-u-fail (+ pl-u-fail 1))
          (set!
            pl-u-failures
            (cons (list name :expected expected :got got) pl-u-failures)))))))

;; Shortcuts
(define pl-a (fn (n) (list "atom" n)))
(define pl-n (fn (v) (list "num" v)))
(define pl-s (fn (v) (list "str" v)))
(define pl-c (fn (f args) (list "compound" f args)))

;; ── Primitive predicates ──────────────────────────────────────────
(pl-u-test! "var? on fresh var" (fn () (pl-var? (pl-mk-rt-var "X"))) true)

(pl-u-test! "var? on atom" (fn () (pl-var? (pl-a "foo"))) false)

(pl-u-test! "atom? on atom" (fn () (pl-atom? (pl-a "foo"))) true)

(pl-u-test! "atom? on var" (fn () (pl-atom? (pl-mk-rt-var "X"))) false)

(pl-u-test!
  "compound? on compound"
  (fn () (pl-compound? (pl-c "p" (list (pl-a "a")))))
  true)

(pl-u-test! "num? on num" (fn () (pl-num? (pl-n 42))) true)

;; ── Fresh var ids ─────────────────────────────────────────────────
(pl-u-test!
  "fresh vars get distinct ids"
  (fn
    ()
    (let
      ((a (pl-mk-rt-var "X")) (b (pl-mk-rt-var "X")))
      (not (= (pl-var-id a) (pl-var-id b)))))
  true)

;; ── Walk ───────────────────────────────────────────────────────────
(pl-u-test!
  "walk returns unbound var unchanged"
  (fn
    ()
    (let
      ((v (pl-mk-rt-var "X")))
      (= (pl-var-id (pl-walk v)) (pl-var-id v))))
  true)

(pl-u-test!
  "walk follows single binding"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (v (pl-mk-rt-var "X")))
      (pl-bind! v (pl-a "hello") t)
      (pl-walk v)))
  (list "atom" "hello"))

(pl-u-test!
  "walk follows chain var→var→atom"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (a (pl-mk-rt-var "A")) (b (pl-mk-rt-var "B")))
      (pl-bind! a b t)
      (pl-bind! b (pl-a "end") t)
      (pl-walk a)))
  (list "atom" "end"))

;; ── Unify: atoms ──────────────────────────────────────────────────
(pl-u-test!
  "unify same atom"
  (fn () (pl-unify! (pl-a "foo") (pl-a "foo") (pl-mk-trail)))
  true)

(pl-u-test!
  "unify different atoms"
  (fn () (pl-unify! (pl-a "foo") (pl-a "bar") (pl-mk-trail)))
  false)

;; ── Unify: numbers ────────────────────────────────────────────────
(pl-u-test!
  "unify equal nums"
  (fn () (pl-unify! (pl-n 5) (pl-n 5) (pl-mk-trail)))
  true)

(pl-u-test!
  "unify different nums"
  (fn () (pl-unify! (pl-n 5) (pl-n 6) (pl-mk-trail)))
  false)

(pl-u-test!
  "atom vs num fails"
  (fn () (pl-unify! (pl-a "5") (pl-n 5) (pl-mk-trail)))
  false)

;; ── Unify: strings ────────────────────────────────────────────────
(pl-u-test!
  "unify equal strings"
  (fn () (pl-unify! (pl-s "hi") (pl-s "hi") (pl-mk-trail)))
  true)

(pl-u-test!
  "unify different strings"
  (fn () (pl-unify! (pl-s "hi") (pl-s "bye") (pl-mk-trail)))
  false)

;; ── Unify: variables ──────────────────────────────────────────────
(pl-u-test!
  "unify var with atom binds"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify! x (pl-a "foo") t)
      (pl-walk x)))
  (list "atom" "foo"))

(pl-u-test!
  "unify atom with var binds"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify! (pl-a "foo") x t)
      (pl-walk x)))
  (list "atom" "foo"))

(pl-u-test!
  "unify var = var binds one to the other"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-unify! x y t)
      (pl-unify! y (pl-a "bound") t)
      (pl-walk x)))
  (list "atom" "bound"))

(pl-u-test!
  "unify same var succeeds without binding"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify! x x t)
      (list (pl-var-bound? x) (= (dict-get t :len) 0))))
  (list false true))

(pl-u-test!
  "bound-var vs atom uses binding"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-bind! x (pl-a "a") t)
      (pl-unify! x (pl-a "a") t)))
  true)

(pl-u-test!
  "bound-var vs different atom fails"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-bind! x (pl-a "a") t)
      (pl-unify! x (pl-a "b") t)))
  false)

;; ── Unify: compounds ──────────────────────────────────────────────
(pl-u-test!
  "unify p(a) with p(a)"
  (fn
    ()
    (pl-unify!
      (pl-c "p" (list (pl-a "a")))
      (pl-c "p" (list (pl-a "a")))
      (pl-mk-trail)))
  true)

(pl-u-test!
  "unify p(a) with p(b)"
  (fn
    ()
    (pl-unify!
      (pl-c "p" (list (pl-a "a")))
      (pl-c "p" (list (pl-a "b")))
      (pl-mk-trail)))
  false)

(pl-u-test!
  "unify p(a) with q(a) — functor mismatch"
  (fn
    ()
    (pl-unify!
      (pl-c "p" (list (pl-a "a")))
      (pl-c "q" (list (pl-a "a")))
      (pl-mk-trail)))
  false)

(pl-u-test!
  "unify p(a) with p(a,b) — arity mismatch"
  (fn
    ()
    (pl-unify!
      (pl-c "p" (list (pl-a "a")))
      (pl-c "p" (list (pl-a "a") (pl-a "b")))
      (pl-mk-trail)))
  false)

(pl-u-test!
  "unify p(X) with p(foo) binds X"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify! (pl-c "p" (list x)) (pl-c "p" (list (pl-a "foo"))) t)
      (pl-walk x)))
  (list "atom" "foo"))

(pl-u-test!
  "unify p(X,Y) with p(1,2) binds both"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-unify!
        (pl-c "p" (list x y))
        (pl-c "p" (list (pl-n 1) (pl-n 2)))
        t)
      (list (pl-walk x) (pl-walk y))))
  (list (list "num" 1) (list "num" 2)))

(pl-u-test!
  "unify p(X,X) with p(a,a)"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify!
        (pl-c "p" (list x x))
        (pl-c "p" (list (pl-a "a") (pl-a "a")))
        t)))
  true)

(pl-u-test!
  "unify p(X,X) with p(a,b) fails"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify!
        (pl-c "p" (list x x))
        (pl-c "p" (list (pl-a "a") (pl-a "b")))
        t)))
  false)

;; ── Nested compounds ──────────────────────────────────────────────
(pl-u-test!
  "unify f(g(X)) with f(g(foo))"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify!
        (pl-c "f" (list (pl-c "g" (list x))))
        (pl-c "f" (list (pl-c "g" (list (pl-a "foo")))))
        t)
      (pl-walk x)))
  (list "atom" "foo"))

;; ── Cons-cell lists ──────────────────────────────────────────────
(define pl-nil (pl-a "[]"))
(define pl-cons (fn (h t) (pl-c "." (list h t))))

(pl-u-test!
  "unify [1,2,3] with [X|T] binds X and T"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (tl (pl-mk-rt-var "T")))
      (pl-unify!
        (pl-cons (pl-n 1) (pl-cons (pl-n 2) (pl-cons (pl-n 3) pl-nil)))
        (pl-cons x tl)
        t)
      (list (pl-walk x) (pl-walk-deep tl))))
  (list
    (list "num" 1)
    (list
      "compound"
      "."
      (list
        (list "num" 2)
        (list "compound" "." (list (list "num" 3) (list "atom" "[]")))))))

;; ── Trail: mark + undo ───────────────────────────────────────────
(pl-u-test!
  "trail-undo restores unbound vars"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (let
        ((mark (pl-trail-mark t)))
        (pl-unify! x (pl-a "a") t)
        (pl-unify! y (pl-a "b") t)
        (pl-trail-undo-to! t mark)
        (list (pl-var-bound? x) (pl-var-bound? y) (dict-get t :len)))))
  (list false false 0))

(pl-u-test!
  "partial undo preserves earlier bindings"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-unify! x (pl-a "a") t)
      (let
        ((mark (pl-trail-mark t)))
        (pl-unify! y (pl-a "b") t)
        (pl-trail-undo-to! t mark)
        (list (pl-var-bound? x) (pl-var-bound? y)))))
  (list true false))

(pl-u-test!
  "try-unify undoes on failure"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-try-unify!
        (pl-c "p" (list x y (pl-a "a")))
        (pl-c "p" (list (pl-n 1) (pl-n 2) (pl-a "b")))
        t)
      (list (pl-var-bound? x) (pl-var-bound? y))))
  (list false false))

(pl-u-test!
  "try-unify success keeps bindings"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-try-unify!
        (pl-c "p" (list x y))
        (pl-c "p" (list (pl-n 1) (pl-n 2)))
        t)
      (list (pl-walk x) (pl-walk y))))
  (list (list "num" 1) (list "num" 2)))

;; ── Occurs check ──────────────────────────────────────────────────
(pl-u-test!
  "no occurs check: X = f(X) succeeds"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (pl-unify! x (pl-c "f" (list x)) t)))
  true)

(pl-u-test!
  "occurs check: X = f(X) fails"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (dict-set! t :occurs-check true)
      (pl-unify! x (pl-c "f" (list x)) t)))
  false)

(pl-u-test!
  "occurs check: X = Y, Y = f(X) fails"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (dict-set! t :occurs-check true)
      (pl-unify! x y t)
      (pl-unify! y (pl-c "f" (list x)) t)))
  false)

(pl-u-test!
  "occurs check: deep occurrence in compound"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")))
      (dict-set! t :occurs-check true)
      (pl-unify! x (pl-c "f" (list (pl-c "g" (list x)))) t)))
  false)

(pl-u-test!
  "occurs check: X = Y (both unbound) succeeds"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (dict-set! t :occurs-check true)
      (pl-unify! x y t)))
  true)

;; ── Parse AST → runtime term ──────────────────────────────────────
(pl-u-test!
  "instantiate replaces (var X) with runtime var"
  (fn
    ()
    (let ((ast (list "var" "X"))) (pl-var? (pl-instantiate-fresh ast))))
  true)

(pl-u-test!
  "instantiate shares vars within clause"
  (fn
    ()
    (let
      ((env {}))
      (let
        ((v1 (pl-instantiate (list "var" "X") env))
          (v2 (pl-instantiate (list "var" "X") env)))
        (= (pl-var-id v1) (pl-var-id v2)))))
  true)

(pl-u-test!
  "instantiate makes distinct vars for _"
  (fn
    ()
    (let
      ((env {}))
      (let
        ((v1 (pl-instantiate (list "var" "_") env))
          (v2 (pl-instantiate (list "var" "_") env)))
        (not (= (pl-var-id v1) (pl-var-id v2))))))
  true)

(pl-u-test!
  "instantiate compound recurses"
  (fn
    ()
    (let
      ((env {}))
      (let
        ((inst (pl-instantiate (list "compound" "p" (list (list "var" "X") (list "atom" "a"))) env))
          (x (dict-get env "X")))
        (pl-unify!
          inst
          (pl-c "p" (list (pl-a "foo") (pl-a "a")))
          (pl-mk-trail))
        (pl-walk x))))
  (list "atom" "foo"))

(pl-u-test!
  "deep-walk resolves nested vars"
  (fn
    ()
    (let
      ((t (pl-mk-trail)) (x (pl-mk-rt-var "X")) (y (pl-mk-rt-var "Y")))
      (pl-unify! x (pl-c "f" (list y (pl-a "b"))) t)
      (pl-unify! y (pl-a "a") t)
      (pl-walk-deep x)))
  (list "compound" "f" (list (list "atom" "a") (list "atom" "b"))))

;; ── Runner ────────────────────────────────────────────────────────
(define pl-unify-tests-run! (fn () {:failed pl-u-fail :passed pl-u-pass :total pl-u-count :failures pl-u-failures}))