;; ==========================================================================
;; z3.sx — SX spec to SMT-LIB translator, written in SX
;;
;; Translates define-primitive, define-io-primitive, and define-special-form
;; declarations from the SX spec into SMT-LIB verification conditions for
;; Z3 and other theorem provers.
;;
;; This is the first self-hosted bootstrapper: the SX evaluator (itself
;; bootstrapped from eval.sx) executes this file against the spec to
;; produce output in a different language. Same pattern as bootstrap_js.py
;; and bootstrap_py.py, but written in SX instead of Python.
;;
;; Usage (from SX):
;;   (z3-translate expr)         — translate one define-* form
;;   (z3-translate-file exprs)   — translate a list of parsed expressions
;;
;; Usage (as reader macro):
;;   #z3(define-primitive "inc" :params (n) :returns "number" :body (+ n 1))
;;   → "; inc — ...\n(declare-fun inc (Int) Int)\n..."
;; ==========================================================================


;; --------------------------------------------------------------------------
;; Type mapping: SX type names → SMT-LIB sorts
;; --------------------------------------------------------------------------

(define z3-sort
  (fn ((sx-type :as string))
    (case sx-type
      "number"  "Int"
      "boolean" "Bool"
      "string"  "String"
      "list"    "(List Value)"
      "dict"    "(Array String Value)"
      :else     "Value")))


;; --------------------------------------------------------------------------
;; Name translation: SX identifiers → SMT-LIB identifiers
;; --------------------------------------------------------------------------

(define z3-name
  (fn ((name :as string))
    (cond
      (= name "!=") "neq"
      (= name "+")  "+"
      (= name "-")  "-"
      (= name "*")  "*"
      (= name "/")  "/"
      (= name "=")  "="
      (= name "<")  "<"
      (= name ">")  ">"
      (= name "<=") "<="
      (= name ">=") ">="
      :else (replace (replace (replace name "-" "_") "?" "_p") "!" "_bang"))))

(define z3-sym
  (fn (sym)
    (let ((name (symbol-name sym)))
      (cond
        (ends-with? name "?")
          (str "is_" (replace (slice name 0 (- (string-length name) 1)) "-" "_"))
        :else
          (replace (replace name "-" "_") "!" "_bang")))))


;; --------------------------------------------------------------------------
;; Expression translation: SX body expressions → SMT-LIB s-expressions
;; --------------------------------------------------------------------------

;; Operators that pass through unchanged
(define z3-identity-ops
  (list "+" "-" "*" "/" "=" "!=" "<" ">" "<=" ">=" "and" "or" "not" "mod"))

;; Operators that get renamed
(define z3-rename-op
  (fn ((op :as string))
    (case op
      "if"  "ite"
      "str" "str.++"
      :else nil)))

(define z3-expr
  (fn (expr)
    (cond
      ;; Numbers
      (number? expr)
        (str expr)

      ;; Strings
      (string? expr)
        (str "\"" expr "\"")

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

      ;; Nil
      (nil? expr)
        "nil_val"

      ;; Symbols
      (= (type-of expr) "symbol")
        (z3-sym expr)

      ;; Lists (function calls / special forms)
      (list? expr)
        (if (empty? expr)
          "()"
          (let ((head (first expr))
                (args (rest expr)))
            (if (not (= (type-of head) "symbol"))
              (str expr)
              (let ((op (symbol-name head)))
                (cond
                  ;; Identity ops: same syntax in both languages
                  (some (fn (x) (= x op)) z3-identity-ops)
                    (str "(" op " " (join " " (map z3-expr args)) ")")

                  ;; Renamed ops
                  (not (nil? (z3-rename-op op)))
                    (str "(" (z3-rename-op op) " " (join " " (map z3-expr args)) ")")

                  ;; max → ite
                  (and (= op "max") (= (len args) 2))
                    (let ((a (z3-expr (nth args 0)))
                          (b (z3-expr (nth args 1))))
                      (str "(ite (>= " a " " b ") " a " " b ")"))

                  ;; min → ite
                  (and (= op "min") (= (len args) 2))
                    (let ((a (z3-expr (nth args 0)))
                          (b (z3-expr (nth args 1))))
                      (str "(ite (<= " a " " b ") " a " " b ")"))

                  ;; empty? → length check
                  (= op "empty?")
                    (str "(= (len " (z3-expr (first args)) ") 0)")

                  ;; first/rest → list ops
                  (= op "first")
                    (str "(head " (z3-expr (first args)) ")")
                  (= op "rest")
                    (str "(tail " (z3-expr (first args)) ")")

                  ;; reduce with initial value
                  (and (= op "reduce") (>= (len args) 3))
                    (str "(reduce " (z3-expr (nth args 0)) " "
                         (z3-expr (nth args 2)) " "
                         (z3-expr (nth args 1)) ")")

                  ;; fn (lambda)
                  (= op "fn")
                    (let ((params (first args))
                          (body (nth args 1)))
                      (str "(lambda (("
                           (join " " (map (fn (p) (str "(" (z3-sym p) " Int)")) params))
                           ")) " (z3-expr body) ")"))

                  ;; native-* → strip prefix
                  (starts-with? op "native-")
                    (str "(" (slice op 7 (string-length op)) " "
                         (join " " (map z3-expr args)) ")")

                  ;; Generic function call
                  :else
                    (str "(" (z3-name op) " "
                         (join " " (map z3-expr args)) ")"))))))

      ;; Fallback
      :else (str expr))))


;; --------------------------------------------------------------------------
;; Keyword argument extraction from define-* forms
;; --------------------------------------------------------------------------

(define z3-extract-kwargs
  (fn ((expr :as list))
    ;; Returns a dict of keyword args from a define-* form
    ;; (define-primitive "name" :params (...) :returns "type" ...) → {:params ... :returns ...}
    (let ((result {})
          (items (rest (rest expr))))  ;; skip head and name
      (z3-extract-kwargs-loop items result))))

(define z3-extract-kwargs-loop
  (fn ((items :as list) (result :as dict))
    (if (or (empty? items) (< (len items) 2))
      result
      (if (= (type-of (first items)) "keyword")
        (z3-extract-kwargs-loop
          (rest (rest items))
          (assoc result (keyword-name (first items)) (nth items 1)))
        (z3-extract-kwargs-loop (rest items) result)))))


;; --------------------------------------------------------------------------
;; Parameter processing
;; --------------------------------------------------------------------------

(define z3-params-to-sorts
  (fn ((params :as list))
    ;; Convert SX param list to list of (name sort) pairs, skipping &rest/&key
    (z3-params-loop params false (list))))

(define z3-params-loop
  (fn ((params :as list) (skip-next :as boolean) (acc :as list))
    (if (empty? params)
      acc
      (let ((p (first params))
            (rest-p (rest params)))
        (cond
          ;; &rest or &key marker — skip it and the next param
          (and (= (type-of p) "symbol")
               (or (= (symbol-name p) "&rest")
                   (= (symbol-name p) "&key")))
            (z3-params-loop rest-p true acc)
          ;; Skipping the param after &rest/&key
          skip-next
            (z3-params-loop rest-p false acc)
          ;; Normal parameter
          (= (type-of p) "symbol")
            (z3-params-loop rest-p false
              (append acc (list (list (symbol-name p) "Int"))))
          ;; Something else — skip
          :else
            (z3-params-loop rest-p false acc))))))

(define z3-has-rest?
  (fn ((params :as list))
    (some (fn (p) (and (= (type-of p) "symbol") (= (symbol-name p) "&rest")))
          params)))


;; --------------------------------------------------------------------------
;; define-primitive → SMT-LIB
;; --------------------------------------------------------------------------

(define z3-translate-primitive
  (fn ((expr :as list))
    (let ((name      (nth expr 1))
          (kwargs    (z3-extract-kwargs expr))
          (params    (or (get kwargs "params") (list)))
          (returns   (or (get kwargs "returns") "any"))
          (doc       (or (get kwargs "doc") ""))
          (body      (get kwargs "body"))
          (pairs     (z3-params-to-sorts params))
          (has-rest  (z3-has-rest? params))
          (smt-name  (z3-name name)))

      (str
        ;; Comment header
        "; " name " — " doc "\n"

        ;; Declaration
        (if has-rest
          (str "; (variadic — modeled as uninterpreted)\n"
               "(declare-fun " smt-name " (Int Int) " (z3-sort returns) ")")
          (str "(declare-fun " smt-name " ("
               (join " " (map (fn (pair) (nth pair 1)) pairs))
               ") " (z3-sort returns) ")"))
        "\n"

        ;; Assertion (if body exists and not variadic)
        (if (and (not (nil? body)) (not has-rest))
          (if (empty? pairs)
            ;; No params — simple assertion
            (str "(assert (= (" smt-name ") " (z3-expr body) "))\n")
            ;; With params — forall
            (let ((bindings (join " " (map (fn (pair) (str "(" (nth pair 0) " Int)")) pairs)))
                  (call-args (join " " (map (fn (pair) (nth pair 0)) pairs))))
              (str "(assert (forall ((" bindings "))\n"
                   "  (= (" smt-name " " call-args ") " (z3-expr body) ")))\n")))
          "")

        ;; Check satisfiability
        "(check-sat)"))))


;; --------------------------------------------------------------------------
;; define-io-primitive → SMT-LIB
;; --------------------------------------------------------------------------

(define z3-translate-io
  (fn ((expr :as list))
    (let ((name   (nth expr 1))
          (kwargs (z3-extract-kwargs expr))
          (doc    (or (get kwargs "doc") ""))
          (smt-name (replace (replace name "-" "_") "?" "_p")))
      (str "; IO primitive: " name " — " doc "\n"
           "; (uninterpreted — IO cannot be verified statically)\n"
           "(declare-fun " smt-name " () Value)"))))


;; --------------------------------------------------------------------------
;; define-special-form → SMT-LIB
;; --------------------------------------------------------------------------

(define z3-translate-special-form
  (fn ((expr :as list))
    (let ((name   (nth expr 1))
          (kwargs (z3-extract-kwargs expr))
          (doc    (or (get kwargs "doc") "")))
      (case name
        "if"
          (str "; Special form: if — " doc "\n"
               "(assert (forall ((c Bool) (t Value) (e Value))\n"
               "  (= (sx_if c t e) (ite c t e))))\n"
               "(check-sat)")
        "when"
          (str "; Special form: when — " doc "\n"
               "(assert (forall ((c Bool) (body Value))\n"
               "  (= (sx_when c body) (ite c body nil_val))))\n"
               "(check-sat)")
        :else
          (str "; Special form: " name " — " doc "\n"
               "; (not directly expressible in SMT-LIB)")))))


;; --------------------------------------------------------------------------
;; Top-level dispatch
;; --------------------------------------------------------------------------

(define z3-translate
  (fn (expr)
    (if (not (list? expr))
      "; Cannot translate: not a list form"
      (if (< (len expr) 2)
        "; Cannot translate: too short"
        (let ((head (first expr)))
          (if (not (= (type-of head) "symbol"))
            "; Cannot translate: head is not a symbol"
            (case (symbol-name head)
              "define-primitive"    (z3-translate-primitive expr)
              "define-io-primitive" (z3-translate-io expr)
              "define-special-form" (z3-translate-special-form expr)
              :else                (z3-expr expr))))))))


;; --------------------------------------------------------------------------
;; Batch translation: process a list of parsed expressions
;; --------------------------------------------------------------------------

(define z3-translate-file
  (fn ((exprs :as list))
    ;; Filter to translatable forms and translate each
    (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)))
      (join "\n\n" (map z3-translate translatable)))))