;; lib/datalog/db.sx — Datalog database (EDB + IDB + rules) + safety hook.
;;
;; A db is a mutable dict:
;;   {:facts {<rel-name-string> -> (literal ...)}
;;    :rules ({:head literal :body (literal ...)} ...)}
;;
;; Facts are stored as full literals `(rel arg ... arg)` so they unify
;; directly against rule body literals. Each relation's tuple list is
;; deduplicated on insert.
;;
;; Phase 3 introduced safety analysis for head variables; Phase 4 (in
;; lib/datalog/builtins.sx) swaps in the real `dl-rule-check-safety`,
;; which is order-aware and understands built-in predicates.

(define dl-make-db (fn () {:facts {} :rules (list)}))

(define
  dl-rel-name
  (fn
    (lit)
    (cond
      ((and (dict? lit) (has-key? lit :neg)) (dl-rel-name (get lit :neg)))
      ((and (list? lit) (> (len lit) 0) (symbol? (first lit)))
        (symbol->string (first lit)))
      (else nil))))

(define dl-builtin-rels (list "<" "<=" ">" ">=" "=" "!=" "is"))

(define
  dl-member-string?
  (fn
    (s xs)
    (cond
      ((= (len xs) 0) false)
      ((= (first xs) s) true)
      (else (dl-member-string? s (rest xs))))))

(define
  dl-builtin?
  (fn
    (lit)
    (and
      (list? lit)
      (> (len lit) 0)
      (let
        ((rel (dl-rel-name lit)))
        (cond
          ((nil? rel) false)
          (else (dl-member-string? rel dl-builtin-rels)))))))

(define
  dl-positive-lit?
  (fn
    (lit)
    (cond
      ((and (dict? lit) (has-key? lit :neg)) false)
      ((dl-builtin? lit) false)
      ((and (list? lit) (> (len lit) 0)) true)
      (else false))))

(define
  dl-tuple-equal?
  (fn
    (a b)
    (cond
      ((and (list? a) (list? b))
        (and (= (len a) (len b)) (dl-tuple-equal-list? a b 0)))
      ((and (number? a) (number? b)) (= a b))
      (else (equal? a b)))))

(define
  dl-tuple-equal-list?
  (fn
    (a b i)
    (cond
      ((>= i (len a)) true)
      ((not (dl-tuple-equal? (nth a i) (nth b i))) false)
      (else (dl-tuple-equal-list? a b (+ i 1))))))

(define
  dl-tuple-member?
  (fn
    (lit lits)
    (cond
      ((= (len lits) 0) false)
      ((dl-tuple-equal? lit (first lits)) true)
      (else (dl-tuple-member? lit (rest lits))))))

(define
  dl-ensure-rel!
  (fn
    (db rel-key)
    (let
      ((facts (get db :facts)))
      (do
        (when
          (not (has-key? facts rel-key))
          (dict-set! facts rel-key (list)))
        (get facts rel-key)))))

(define
  dl-rel-tuples
  (fn
    (db rel-key)
    (let
      ((facts (get db :facts)))
      (if (has-key? facts rel-key) (get facts rel-key) (list)))))

(define
  dl-add-fact!
  (fn
    (db lit)
    (cond
      ((not (and (list? lit) (> (len lit) 0)))
        (error (str "dl-add-fact!: expected literal list, got " lit)))
      ((not (dl-ground? lit (dl-empty-subst)))
        (error (str "dl-add-fact!: expected ground literal, got " lit)))
      (else
        (let
          ((rel-key (dl-rel-name lit)))
          (let
            ((tuples (dl-ensure-rel! db rel-key)))
            (cond
              ((dl-tuple-member? lit tuples) false)
              (else (do (append! tuples lit) true)))))))))

;; The full safety check lives in builtins.sx (it has to know which
;; predicates are built-ins). dl-add-rule! calls it via forward
;; reference; load builtins.sx alongside db.sx in any setup that
;; adds rules. The fallback below is used if builtins.sx isn't loaded.
(define
  dl-rule-check-safety
  (fn
    (rule)
    (let
      ((head-vars (dl-vars-of (get rule :head))) (body-vars (list)))
      (do
        (for-each
          (fn
            (lit)
            (when
              (and
                (list? lit)
                (> (len lit) 0)
                (not (and (dict? lit) (has-key? lit :neg))))
              (for-each
                (fn
                  (v)
                  (when
                    (not (dl-member-string? v body-vars))
                    (append! body-vars v)))
                (dl-vars-of lit))))
          (get rule :body))
        (let
          ((missing (list)))
          (do
            (for-each
              (fn
                (v)
                (when
                  (and
                    (not (dl-member-string? v body-vars))
                    (not (= v "_")))
                  (append! missing v)))
              head-vars)
            (cond
              ((> (len missing) 0)
                (str
                  "head variable(s) "
                  missing
                  " do not appear in any body literal"))
              (else nil))))))))

(define
  dl-add-rule!
  (fn
    (db rule)
    (cond
      ((not (dict? rule))
        (error (str "dl-add-rule!: expected rule dict, got " rule)))
      ((not (has-key? rule :head))
        (error (str "dl-add-rule!: rule missing :head, got " rule)))
      (else
        (let
          ((err (dl-rule-check-safety rule)))
          (cond
            ((not (nil? err)) (error (str "dl-add-rule!: " err)))
            (else
              (let
                ((rules (get db :rules)))
                (do (append! rules rule) true)))))))))

(define
  dl-add-clause!
  (fn
    (db clause)
    (cond
      ((has-key? clause :query) false)
      ((and (has-key? clause :body) (= (len (get clause :body)) 0))
        (dl-add-fact! db (get clause :head)))
      (else (dl-add-rule! db clause)))))

(define
  dl-load-program!
  (fn
    (db source)
    (let
      ((clauses (dl-parse source)))
      (do (for-each (fn (c) (dl-add-clause! db c)) clauses) db))))

(define
  dl-program
  (fn (source) (let ((db (dl-make-db))) (dl-load-program! db source))))

(define dl-rules (fn (db) (get db :rules)))

(define
  dl-fact-count
  (fn
    (db)
    (let
      ((facts (get db :facts)) (total 0))
      (do
        (for-each
          (fn (k) (set! total (+ total (len (get facts k)))))
          (keys facts))
        total))))