;; 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 {} :facts-keys {} :facts-index {} :rules (list) :strategy :semi-naive}))

;; Evaluation strategy. Default :semi-naive (the only strategy
;; currently implemented). :magic is reserved for goal-directed
;; magic-sets evaluation — calling it now logs a one-time "deferred"
;; note and falls back to semi-naive.
(define
  dl-set-strategy!
  (fn
    (db strategy)
    (do
      (dict-set! db :strategy strategy)
      db)))

(define
  dl-get-strategy
  (fn
    (db)
    (if (has-key? db :strategy) (get db :strategy) :semi-naive)))

(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)
    (dl-tuple-member-aux? lit lits 0 (len lits))))

(define
  dl-tuple-member-aux?
  (fn
    (lit lits i n)
    (cond
      ((>= i n) false)
      ((dl-tuple-equal? lit (nth lits i)) true)
      (else (dl-tuple-member-aux? lit lits (+ i 1) n)))))

(define
  dl-ensure-rel!
  (fn
    (db rel-key)
    (let
      ((facts (get db :facts))
        (fk (get db :facts-keys))
        (fi (get db :facts-index)))
      (do
        (when
          (not (has-key? facts rel-key))
          (dict-set! facts rel-key (list)))
        (when
          (not (has-key? fk rel-key))
          (dict-set! fk rel-key {}))
        (when
          (not (has-key? fi rel-key))
          (dict-set! fi rel-key {}))
        (get facts rel-key)))))

;; First-arg index helpers. Tuples are keyed by their first-after-rel
;; arg's `(str ...)`; when that arg is a constant, dl-match-positive
;; uses the index instead of scanning the full relation.
(define
  dl-arg-key
  (fn
    (v)
    (str v)))

(define
  dl-index-add!
  (fn
    (db rel-key lit)
    (let
      ((idx (get db :facts-index))
        (n (len lit)))
      (when
        (and (>= n 2) (has-key? idx rel-key))
        (let
          ((rel-idx (get idx rel-key))
            (k (dl-arg-key (nth lit 1))))
          (do
            (when
              (not (has-key? rel-idx k))
              (dict-set! rel-idx k (list)))
            (append! (get rel-idx k) lit)))))))

(define
  dl-index-lookup
  (fn
    (db rel-key arg-val)
    (let
      ((idx (get db :facts-index)))
      (cond
        ((not (has-key? idx rel-key)) (list))
        (else
          (let ((rel-idx (get idx rel-key))
                (k (dl-arg-key arg-val)))
            (if (has-key? rel-idx k) (get rel-idx k) (list))))))))

(define dl-tuple-key (fn (lit) (str lit)))

(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))
              (key-dict (get (get db :facts-keys) rel-key))
              (tk (dl-tuple-key lit)))
            (cond
              ((has-key? key-dict tk) false)
              (else
                (do
                  (dict-set! key-dict tk true)
                  (append! tuples lit)
                  (dl-index-add! db rel-key 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-rename-anon-term
  (fn
    (term next-name)
    (cond
      ((and (symbol? term) (= (symbol->string term) "_"))
        (next-name))
      ((list? term)
        (map (fn (x) (dl-rename-anon-term x next-name)) term))
      (else term))))

(define
  dl-rename-anon-lit
  (fn
    (lit next-name)
    (cond
      ((and (dict? lit) (has-key? lit :neg))
        {:neg (dl-rename-anon-term (get lit :neg) next-name)})
      ((list? lit) (dl-rename-anon-term lit next-name))
      (else lit))))

(define
  dl-make-anon-renamer
  (fn
    ()
    (let ((counter 0))
      (fn () (do (set! counter (+ counter 1))
                 (string->symbol (str "_anon" counter)))))))

(define
  dl-rename-anon-rule
  (fn
    (rule)
    (let ((next-name (dl-make-anon-renamer)))
      {:head (dl-rename-anon-term (get rule :head) next-name)
       :body (map (fn (lit) (dl-rename-anon-lit lit next-name))
                  (get rule :body))})))

(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 ((rule (dl-rename-anon-rule rule)))
          (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))))

;; Returns {<rel-name>: tuple-count} for debugging. Includes
;; relations with any tuples plus all rule-head relations (so empty
;; IDB shows as 0). Skips empty EDB-only entries that are placeholders
;; from internal `dl-ensure-rel!` calls.
(define
  dl-summary
  (fn
    (db)
    (let
      ((facts (get db :facts))
        (out {})
        (rule-heads (list)))
      (do
        (for-each
          (fn
            (rule)
            (let ((h (dl-rel-name (get rule :head))))
              (when
                (and (not (nil? h)) (not (dl-member-string? h rule-heads)))
                (append! rule-heads h))))
          (dl-rules db))
        (for-each
          (fn
            (k)
            (let ((c (len (get facts k))))
              (when
                (or (> c 0) (dl-member-string? k rule-heads))
                (dict-set! out k c))))
          (keys facts))
        ;; Add rule heads that have no facts (yet).
        (for-each
          (fn
            (k)
            (when (not (has-key? out k)) (dict-set! out k 0)))
          rule-heads)
        out))))