;; ==========================================================================
;; prove.sx — SMT-LIB satisfiability checker, written in SX
;;
;; Verifies the SMT-LIB output from z3.sx. For the class of assertions
;; z3.sx produces (definitional equalities), satisfiability is provable
;; by construction: the definition IS the model.
;;
;; This closes the loop:
;;   primitives.sx → z3.sx → SMT-LIB → prove.sx → sat
;;   SX spec → SX translator → s-expressions → SX prover → proof
;;
;; The prover also evaluates each definition with concrete test values
;; to demonstrate consistency.
;;
;; Usage:
;;   (prove-check smtlib-string)     — verify a single check-sat block
;;   (prove-translate expr)          — translate + verify a define-* form
;;   (prove-file exprs)              — verify all define-* forms
;; ==========================================================================


;; --------------------------------------------------------------------------
;; SMT-LIB expression evaluator
;; --------------------------------------------------------------------------

;; Evaluate an SMT-LIB expression in a variable environment
(define smt-eval
  (fn (expr env)
    (cond
      ;; Numbers
      (number? expr) expr

      ;; String literals
      (string? expr)
        (cond
          (= expr "true")  true
          (= expr "false") false
          :else expr)

      ;; Booleans
      (= expr true) true
      (= expr false) false

      ;; Symbols — look up in env
      (= (type-of expr) "symbol")
        (let ((name (symbol-name expr)))
          (cond
            (= name "true") true
            (= name "false") false
            :else (get env name expr)))

      ;; Lists — function application
      (list? expr)
        (if (empty? expr) nil
          (let ((head (first expr))
                (args (rest expr)))
            (if (not (= (type-of head) "symbol"))
              expr
              (let ((op (symbol-name head)))
                (cond
                  ;; Arithmetic
                  (= op "+")
                    (reduce (fn (a b) (+ a b)) 0
                      (map (fn (a) (smt-eval a env)) args))
                  (= op "-")
                    (if (= (len args) 1)
                      (- 0 (smt-eval (first args) env))
                      (- (smt-eval (nth args 0) env)
                         (smt-eval (nth args 1) env)))
                  (= op "*")
                    (reduce (fn (a b) (* a b)) 1
                      (map (fn (a) (smt-eval a env)) args))
                  (= op "/")
                    (let ((a (smt-eval (nth args 0) env))
                          (b (smt-eval (nth args 1) env)))
                      (if (= b 0) 0 (/ a b)))
                  (= op "div")
                    (let ((a (smt-eval (nth args 0) env))
                          (b (smt-eval (nth args 1) env)))
                      (if (= b 0) 0 (/ a b)))
                  (= op "mod")
                    (let ((a (smt-eval (nth args 0) env))
                          (b (smt-eval (nth args 1) env)))
                      (if (= b 0) 0 (mod a b)))

                  ;; Comparison
                  (= op "=")
                    (= (smt-eval (nth args 0) env)
                       (smt-eval (nth args 1) env))
                  (= op "<")
                    (< (smt-eval (nth args 0) env)
                       (smt-eval (nth args 1) env))
                  (= op ">")
                    (> (smt-eval (nth args 0) env)
                       (smt-eval (nth args 1) env))
                  (= op "<=")
                    (<= (smt-eval (nth args 0) env)
                        (smt-eval (nth args 1) env))
                  (= op ">=")
                    (>= (smt-eval (nth args 0) env)
                        (smt-eval (nth args 1) env))

                  ;; Logic
                  (= op "and")
                    (every? (fn (a) (smt-eval a env)) args)
                  (= op "or")
                    (some (fn (a) (smt-eval a env)) args)
                  (= op "not")
                    (not (smt-eval (first args) env))

                  ;; ite (if-then-else)
                  (= op "ite")
                    (if (smt-eval (nth args 0) env)
                      (smt-eval (nth args 1) env)
                      (smt-eval (nth args 2) env))

                  ;; Function call — look up in env
                  :else
                    (let ((fn-def (get env op nil)))
                      (if (nil? fn-def)
                        (list op (map (fn (a) (smt-eval a env)) args))
                        ;; fn-def is {:params [...] :body expr}
                        (let ((params (get fn-def "params" (list)))
                              (body   (get fn-def "body" nil))
                              (evals  (map (fn (a) (smt-eval a env)) args)))
                          (if (nil? body)
                            ;; Uninterpreted — return symbolic
                            (list op evals)
                            ;; Evaluate body with params bound
                            (smt-eval body
                              (merge env
                                (smt-bind-params params evals))))))))))))

      :else expr)))


;; Bind parameter names to values
(define smt-bind-params
  (fn (params vals)
    (smt-bind-loop params vals {})))

(define smt-bind-loop
  (fn (params vals acc)
    (if (or (empty? params) (empty? vals))
      acc
      (smt-bind-loop (rest params) (rest vals)
        (assoc acc (first params) (first vals))))))


;; --------------------------------------------------------------------------
;; SMT-LIB statement parser
;; --------------------------------------------------------------------------

;; Extract declarations and assertions from parsed SMT-LIB
(define smt-extract-statements
  (fn (exprs)
    (smt-extract-loop exprs {} (list))))

(define smt-extract-loop
  (fn (exprs decls assertions)
    (if (empty? exprs)
      {:decls decls :assertions assertions}
      (let ((expr (first exprs))
            (rest-e (rest exprs)))
        (if (not (list? expr))
          (smt-extract-loop rest-e decls assertions)
          (if (empty? expr)
            (smt-extract-loop rest-e decls assertions)
            (let ((head (symbol-name (first expr))))
              (cond
                ;; (declare-fun name (sorts) sort)
                (= head "declare-fun")
                  (let ((name (nth expr 1))
                        (param-sorts (nth expr 2))
                        (ret-sort (nth expr 3)))
                    (smt-extract-loop rest-e
                      (assoc decls (if (= (type-of name) "symbol")
                                     (symbol-name name) name)
                        {:params (if (list? param-sorts)
                                   (map (fn (s) (if (= (type-of s) "symbol")
                                                  (symbol-name s) (str s)))
                                        param-sorts)
                                   (list))
                         :ret (if (= (type-of ret-sort) "symbol")
                                (symbol-name ret-sort) (str ret-sort))})
                      assertions))

                ;; (assert ...)
                (= head "assert")
                  (smt-extract-loop rest-e decls
                    (append assertions (list (nth expr 1))))

                ;; (check-sat) — skip
                (= head "check-sat")
                  (smt-extract-loop rest-e decls assertions)

                ;; comments (strings starting with ;) — skip
                :else
                  (smt-extract-loop rest-e decls assertions)))))))))


;; --------------------------------------------------------------------------
;; Assertion classifier
;; --------------------------------------------------------------------------

;; Check if an assertion is definitional: (forall (...) (= (f ...) body))
;; or (= (f) body) for nullary
(define smt-definitional?
  (fn (assertion)
    (if (not (list? assertion)) false
      (let ((head (symbol-name (first assertion))))
        (cond
          ;; (forall ((bindings)) (= (f ...) body))
          (= head "forall")
            (let ((body (nth assertion 2)))
              (and (list? body)
                   (= (symbol-name (first body)) "=")))
          ;; (= (f ...) body)
          (= head "=")
            true
          :else false)))))


;; Extract the function name, parameters, and body from a definitional assertion
(define smt-extract-definition
  (fn (assertion)
    (let ((head (symbol-name (first assertion))))
      (cond
        ;; (forall (((x Int) (y Int))) (= (f x y) body))
        (= head "forall")
          (let ((bindings (first (nth assertion 1)))
                (eq-expr (nth assertion 2))
                (call (nth eq-expr 1))
                (body (nth eq-expr 2)))
            {:name (if (= (type-of (first call)) "symbol")
                     (symbol-name (first call)) (str (first call)))
             :params (map (fn (b)
                       (if (list? b)
                         (if (= (type-of (first b)) "symbol")
                           (symbol-name (first b)) (str (first b)))
                         (if (= (type-of b) "symbol")
                           (symbol-name b) (str b))))
                     (if (list? bindings) bindings (list bindings)))
             :body body})

        ;; (= (f) body)
        (= head "=")
          (let ((call (nth assertion 1))
                (body (nth assertion 2)))
            {:name (if (list? call)
                     (if (= (type-of (first call)) "symbol")
                       (symbol-name (first call)) (str (first call)))
                     (str call))
             :params (list)
             :body body})

        :else nil))))


;; --------------------------------------------------------------------------
;; Test value generation
;; --------------------------------------------------------------------------

(define smt-test-values
  (list
    (list 0)
    (list 1)
    (list -1)
    (list 5)
    (list 42)
    (list 1 2)
    (list -3 7)
    (list 5 5)
    (list 100 -50)
    (list 3 1)
    (list 1 1 10)
    (list 5 1 3)
    (list -5 1 10)
    (list 3 3 3)
    (list 7 2 9)))


;; --------------------------------------------------------------------------
;; Proof engine
;; --------------------------------------------------------------------------

;; Verify a single definitional assertion by construction + evaluation
(define smt-verify-definition
  (fn (def-info decls)
    (let ((name (get def-info "name"))
          (params (get def-info "params"))
          (body (get def-info "body"))
          (n-params (len params)))

      ;; Build the model: define f = λparams.body
      (let ((model (assoc decls name {:params params :body body}))
            ;; Select test values matching arity
            (tests (filter (fn (tv) (= (len tv) n-params)) smt-test-values))
            ;; Run tests
            (results (map
              (fn (test-vals)
                (let ((env (merge model (smt-bind-params params test-vals)))
                      ;; Evaluate body directly
                      (body-result (smt-eval body env))
                      ;; Evaluate via function call
                      (call-expr (cons (first (sx-parse name)) test-vals))
                      (call-result (smt-eval call-expr env)))
                  {:vals test-vals
                   :body-result body-result
                   :call-result call-result
                   :equal (= body-result call-result)}))
              tests)))
        {:name name
         :status (if (every? (fn (r) (get r "equal")) results) "sat" "FAIL")
         :proof "by construction (definition is the model)"
         :tests-passed (len (filter (fn (r) (get r "equal")) results))
         :tests-total (len results)
         :sample (if (empty? results) nil (first results))}))))


;; --------------------------------------------------------------------------
;; Public API
;; --------------------------------------------------------------------------

;; Strip SMT-LIB comment lines (starting with ;) and return only actual forms.
;; Handles comments that contain ( characters.
(define smt-strip-comments
  (fn (s)
    (let ((lines (split s "\n"))
          (non-comment (filter
            (fn (line) (not (starts-with? (trim line) ";")))
            lines)))
      (join "\n" non-comment))))

;; Verify SMT-LIB output (string) — parse, classify, prove
(define prove-check
  (fn (smtlib-str)
    (let ((parsed (sx-parse (smt-strip-comments smtlib-str)))
          (stmts (smt-extract-statements parsed))
          (decls (get stmts "decls"))
          (assertions (get stmts "assertions")))
      (if (empty? assertions)
        {:status "sat" :reason "no assertions (declaration only)"}
        (let ((results (map
                (fn (assertion)
                  (if (smt-definitional? assertion)
                    (let ((def-info (smt-extract-definition assertion)))
                      (if (nil? def-info)
                        {:status "unknown" :reason "could not parse definition"}
                        (smt-verify-definition def-info decls)))
                    {:status "unknown"
                     :reason "non-definitional assertion (needs full SMT solver)"}))
                assertions)))
          {:status (if (every? (fn (r) (= (get r "status") "sat")) results)
                    "sat" "unknown")
           :assertions (len assertions)
           :results results})))))


;; Translate a define-* form AND verify it — the full pipeline
(define prove-translate
  (fn (expr)
    (let ((smtlib (z3-translate expr))
          (proof (prove-check smtlib))
          (status (get proof "status"))
          (results (get proof "results" (list))))
      (str smtlib "\n"
           ";; ─── prove.sx ───\n"
           ";; status: " status "\n"
           (if (empty? results) ""
             (let ((r (first results)))
               (str ";; proof:  " (get r "proof" "") "\n"
                    ";; tested: " (str (get r "tests-passed" 0))
                      "/" (str (get r "tests-total" 0))
                      " ground instances\n")))))))


;; Batch verify: translate and prove all define-* forms
(define prove-file
  (fn (exprs)
    (let ((translatable
            (filter
              (fn (expr)
                (and (list? expr)
                     (>= (len expr) 2)
                     (= (type-of (first expr)) "symbol")
                     (let ((name (symbol-name (first expr))))
                       (or (= name "define-primitive")
                           (= name "define-io-primitive")
                           (= name "define-special-form")))))
              exprs))
          (results (map
            (fn (expr)
              (let ((smtlib (z3-translate expr))
                    (proof (prove-check smtlib))
                    (name (nth expr 1)))
                (assoc proof "name" name)))
            translatable))
          (sat-count (len (filter (fn (r) (= (get r "status") "sat")) results)))
          (total (len results)))
      {:total total
       :sat sat-count
       :all-sat (= sat-count total)
       :results results})))